Coverage for python/lsst/daf/butler/registry/dimensions/governor.py : 94%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 (
38 Database,
39 GovernorDimensionRecordStorage,
40 StaticTablesContext,
41)
42from ..queries import QueryBuilder
45class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage):
46 """A record storage implementation for `GovernorDimension` that
47 aggressively fetches and caches all values from the database.
49 Parameters
50 ----------
51 db : `Database`
52 Interface to the database engine and namespace that will hold these
53 dimension records.
54 dimension : `GovernorDimension`
55 The dimension whose records this storage will manage.
56 table : `sqlalchemy.schema.Table`
57 The logical table for the dimension.
58 """
59 def __init__(self, db: Database, dimension: GovernorDimension, table: sqlalchemy.schema.Table):
60 self._db = db
61 self._dimension = dimension
62 self._table = table
63 self._cache: Dict[str, DimensionRecord] = {}
64 self._callbacks: List[Callable[[DimensionRecord], None]] = []
66 @classmethod
67 def initialize(cls, db: Database, element: GovernorDimension, *,
68 context: Optional[StaticTablesContext] = None,
69 config: Mapping[str, Any]) -> GovernorDimensionRecordStorage:
70 # Docstring inherited from GovernorDimensionRecordStorage.
71 spec = element.RecordClass.fields.makeTableSpec(
72 RegionReprClass=db.getSpatialRegionRepresentation(),
73 TimespanReprClass=db.getTimespanRepresentation(),
74 )
75 if context is not None: 75 ↛ 78line 75 didn't jump to line 78, because the condition on line 75 was never false
76 table = context.addTable(element.name, spec)
77 else:
78 table = db.ensureTableExists(element.name, spec)
79 return cls(db, element, table)
81 @property
82 def element(self) -> GovernorDimension:
83 # Docstring inherited from DimensionRecordStorage.element.
84 return self._dimension
86 def refresh(self) -> None:
87 # Docstring inherited from GovernorDimensionRecordStorage.
88 RecordClass = self._dimension.RecordClass
89 sql = sqlalchemy.sql.select(
90 *[self._table.columns[name] for name in RecordClass.fields.standard.names]
91 ).select_from(self._table)
92 cache: Dict[str, DimensionRecord] = {}
93 for row in self._db.query(sql):
94 record = RecordClass(**row._asdict())
95 cache[getattr(record, self._dimension.primaryKey.name)] = record
96 self._cache = cache
98 @property
99 def values(self) -> AbstractSet[str]:
100 # Docstring inherited from GovernorDimensionRecordStorage.
101 return self._cache.keys()
103 @property
104 def table(self) -> sqlalchemy.schema.Table:
105 return self._table
107 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
108 # Docstring inherited from GovernorDimensionRecordStorage.
109 self._callbacks.append(callback)
111 def clearCaches(self) -> None:
112 # Docstring inherited from DimensionRecordStorage.clearCaches.
113 self._cache.clear()
115 def join(
116 self,
117 builder: QueryBuilder, *,
118 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
119 timespans: Optional[NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation]] = None,
120 ) -> None:
121 # Docstring inherited from DimensionRecordStorage.
122 joinOn = builder.startJoin(self._table, self.element.dimensions,
123 self.element.RecordClass.fields.dimensions.names)
124 builder.finishJoin(self._table, joinOn)
125 return self._table
127 def insert(self, *records: DimensionRecord, replace: bool = False) -> None:
128 # Docstring inherited from DimensionRecordStorage.insert.
129 elementRows = [record.toDict() for record in records]
130 with self._db.transaction():
131 if replace: 131 ↛ 132line 131 didn't jump to line 132, because the condition on line 131 was never true
132 self._db.replace(self._table, *elementRows)
133 else:
134 self._db.insert(self._table, *elementRows)
135 for record in records:
136 self._cache[getattr(record, self.element.primaryKey.name)] = record
137 for callback in self._callbacks:
138 callback(record)
140 def sync(self, record: DimensionRecord, update: bool = False) -> Union[bool, Dict[str, Any]]:
141 # Docstring inherited from DimensionRecordStorage.sync.
142 compared = record.toDict()
143 keys = {}
144 for name in record.fields.required.names:
145 keys[name] = compared.pop(name)
146 with self._db.transaction():
147 _, inserted_or_updated = self._db.sync(
148 self._table,
149 keys=keys,
150 compared=compared,
151 update=update,
152 )
153 if inserted_or_updated: 153 ↛ 157line 153 didn't jump to line 157, because the condition on line 153 was never false
154 self._cache[getattr(record, self.element.primaryKey.name)] = record
155 for callback in self._callbacks:
156 callback(record)
157 return inserted_or_updated
159 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
160 # Docstring inherited from DimensionRecordStorage.fetch.
161 try:
162 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
163 except KeyError:
164 pass
165 # If at first we don't succeed, refresh and try again. But this time
166 # we use dict.get to return None if we don't find something.
167 self.refresh()
168 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
170 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
171 # Docstring inherited from DimensionRecordStorage.digestTables.
172 return [self._table]