Coverage for python/lsst/daf/butler/registry/dimensions/governor.py: 96%
81 statements
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-22 02:04 -0700
« prev ^ index » next coverage.py v6.4.4, created at 2022-09-22 02:04 -0700
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
21from __future__ import annotations
23__all__ = ["BasicGovernorDimensionRecordStorage"]
25from typing import AbstractSet, Any, Callable, Dict, Iterable, List, Mapping, Optional, Union
27import sqlalchemy
29from ...core import (
30 DataCoordinateIterable,
31 DimensionElement,
32 DimensionRecord,
33 GovernorDimension,
34 NamedKeyDict,
35 TimespanDatabaseRepresentation,
36)
37from ..interfaces import Database, GovernorDimensionRecordStorage, StaticTablesContext
38from ..queries import QueryBuilder
41class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage):
42 """A record storage implementation for `GovernorDimension` that
43 aggressively fetches and caches all values from the database.
45 Parameters
46 ----------
47 db : `Database`
48 Interface to the database engine and namespace that will hold these
49 dimension records.
50 dimension : `GovernorDimension`
51 The dimension whose records this storage will manage.
52 table : `sqlalchemy.schema.Table`
53 The logical table for the dimension.
54 """
56 def __init__(self, db: Database, dimension: GovernorDimension, table: sqlalchemy.schema.Table):
57 self._db = db
58 self._dimension = dimension
59 self._table = table
60 self._cache: Dict[str, DimensionRecord] = {}
61 self._callbacks: List[Callable[[DimensionRecord], None]] = []
63 @classmethod
64 def initialize(
65 cls,
66 db: Database,
67 element: GovernorDimension,
68 *,
69 context: Optional[StaticTablesContext] = None,
70 config: Mapping[str, Any],
71 ) -> GovernorDimensionRecordStorage:
72 # Docstring inherited from GovernorDimensionRecordStorage.
73 spec = element.RecordClass.fields.makeTableSpec(
74 TimespanReprClass=db.getTimespanRepresentation(),
75 )
76 if context is not None: 76 ↛ 79line 76 didn't jump to line 79, because the condition on line 76 was never false
77 table = context.addTable(element.name, spec)
78 else:
79 table = db.ensureTableExists(element.name, spec)
80 return cls(db, element, table)
82 @property
83 def element(self) -> GovernorDimension:
84 # Docstring inherited from DimensionRecordStorage.element.
85 return self._dimension
87 def refresh(self) -> None:
88 # Docstring inherited from GovernorDimensionRecordStorage.
89 RecordClass = self._dimension.RecordClass
90 sql = sqlalchemy.sql.select(
91 *[self._table.columns[name] for name in RecordClass.fields.standard.names]
92 ).select_from(self._table)
93 cache: Dict[str, DimensionRecord] = {}
94 for row in self._db.query(sql):
95 record = RecordClass(**row._asdict())
96 cache[getattr(record, self._dimension.primaryKey.name)] = record
97 self._cache = cache
99 @property
100 def values(self) -> AbstractSet[str]:
101 # Docstring inherited from GovernorDimensionRecordStorage.
102 return self._cache.keys()
104 @property
105 def table(self) -> sqlalchemy.schema.Table:
106 return self._table
108 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
109 # Docstring inherited from GovernorDimensionRecordStorage.
110 self._callbacks.append(callback)
112 def clearCaches(self) -> None:
113 # Docstring inherited from DimensionRecordStorage.clearCaches.
114 self._cache.clear()
116 def join(
117 self,
118 builder: QueryBuilder,
119 *,
120 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
121 timespans: Optional[NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation]] = None,
122 ) -> None:
123 # Docstring inherited from DimensionRecordStorage.
124 joinOn = builder.startJoin(
125 self._table, self.element.dimensions, self.element.RecordClass.fields.dimensions.names
126 )
127 builder.finishJoin(self._table, joinOn)
128 return self._table
130 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None:
131 # Docstring inherited from DimensionRecordStorage.insert.
132 elementRows = [record.toDict() for record in records]
133 with self._db.transaction():
134 if replace: 134 ↛ 135line 134 didn't jump to line 135, because the condition on line 134 was never true
135 self._db.replace(self._table, *elementRows)
136 elif skip_existing:
137 self._db.ensure(self._table, *elementRows, primary_key_only=True)
138 else:
139 self._db.insert(self._table, *elementRows)
140 for record in records:
141 # We really shouldn't ever get into a situation where the record
142 # here differs from the one in the DB, but the last thing we want
143 # is to make it harder to debug by making the cache different from
144 # the DB.
145 if skip_existing:
146 self._cache.setdefault(getattr(record, self.element.primaryKey.name), record)
147 else:
148 self._cache[getattr(record, self.element.primaryKey.name)] = record
149 for callback in self._callbacks:
150 callback(record)
152 def sync(self, record: DimensionRecord, update: bool = False) -> Union[bool, Dict[str, Any]]:
153 # Docstring inherited from DimensionRecordStorage.sync.
154 compared = record.toDict()
155 keys = {}
156 for name in record.fields.required.names:
157 keys[name] = compared.pop(name)
158 with self._db.transaction():
159 _, inserted_or_updated = self._db.sync(
160 self._table,
161 keys=keys,
162 compared=compared,
163 update=update,
164 )
165 if inserted_or_updated: 165 ↛ 169line 165 didn't jump to line 169, because the condition on line 165 was never false
166 self._cache[getattr(record, self.element.primaryKey.name)] = record
167 for callback in self._callbacks:
168 callback(record)
169 return inserted_or_updated
171 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
172 # Docstring inherited from DimensionRecordStorage.fetch.
173 try:
174 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
175 except KeyError:
176 pass
177 # If at first we don't succeed, refresh and try again. But this time
178 # we use dict.get to return None if we don't find something.
179 self.refresh()
180 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
182 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
183 # Docstring inherited from DimensionRecordStorage.digestTables.
184 return [self._table]