Coverage for python/lsst/daf/butler/registry/dimensions/governor.py: 95%
84 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-26 15:13 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-26 15:13 +0000
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 collections.abc import Callable, Iterable, Mapping, Set
26from typing import Any
28import sqlalchemy
30from ...core import (
31 DataCoordinateIterable,
32 DimensionElement,
33 DimensionRecord,
34 GovernorDimension,
35 NamedKeyDict,
36 TimespanDatabaseRepresentation,
37)
38from ..interfaces import Database, GovernorDimensionRecordStorage, StaticTablesContext
39from ..queries import QueryBuilder
42class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage):
43 """A record storage implementation for `GovernorDimension` that
44 aggressively fetches and caches all values from the database.
46 Parameters
47 ----------
48 db : `Database`
49 Interface to the database engine and namespace that will hold these
50 dimension records.
51 dimension : `GovernorDimension`
52 The dimension whose records this storage will manage.
53 table : `sqlalchemy.schema.Table`
54 The logical table for the dimension.
55 """
57 def __init__(self, db: Database, dimension: GovernorDimension, table: sqlalchemy.schema.Table):
58 self._db = db
59 self._dimension = dimension
60 self._table = table
61 self._cache: dict[str, DimensionRecord] = {}
62 self._callbacks: list[Callable[[DimensionRecord], None]] = []
64 @classmethod
65 def initialize(
66 cls,
67 db: Database,
68 element: GovernorDimension,
69 *,
70 context: StaticTablesContext | None = None,
71 config: Mapping[str, Any],
72 ) -> GovernorDimensionRecordStorage:
73 # Docstring inherited from GovernorDimensionRecordStorage.
74 spec = element.RecordClass.fields.makeTableSpec(
75 TimespanReprClass=db.getTimespanRepresentation(),
76 )
77 if context is not None: 77 ↛ 80line 77 didn't jump to line 80, because the condition on line 77 was never false
78 table = context.addTable(element.name, spec)
79 else:
80 table = db.ensureTableExists(element.name, spec)
81 return cls(db, element, table)
83 @property
84 def element(self) -> GovernorDimension:
85 # Docstring inherited from DimensionRecordStorage.element.
86 return self._dimension
88 def refresh(self) -> None:
89 # Docstring inherited from GovernorDimensionRecordStorage.
90 RecordClass = self._dimension.RecordClass
91 sql = sqlalchemy.sql.select(
92 *[self._table.columns[name] for name in RecordClass.fields.standard.names]
93 ).select_from(self._table)
94 cache: dict[str, DimensionRecord] = {}
95 with self._db.query(sql) as sql_result:
96 sql_rows = sql_result.mappings().fetchall()
97 for row in sql_rows:
98 record = RecordClass(**row)
99 cache[getattr(record, self._dimension.primaryKey.name)] = record
100 self._cache = cache
102 @property
103 def values(self) -> Set[str]:
104 # Docstring inherited from GovernorDimensionRecordStorage.
105 return self._cache.keys()
107 @property
108 def table(self) -> sqlalchemy.schema.Table:
109 return self._table
111 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
112 # Docstring inherited from GovernorDimensionRecordStorage.
113 self._callbacks.append(callback)
115 def clearCaches(self) -> None:
116 # Docstring inherited from DimensionRecordStorage.clearCaches.
117 self._cache.clear()
119 def join(
120 self,
121 builder: QueryBuilder,
122 *,
123 regions: NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement] | None = None,
124 timespans: NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation] | None = None,
125 ) -> None:
126 # Docstring inherited from DimensionRecordStorage.
127 joinOn = builder.startJoin(
128 self._table, self.element.dimensions, self.element.RecordClass.fields.dimensions.names
129 )
130 builder.finishJoin(self._table, joinOn)
131 return self._table
133 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None:
134 # Docstring inherited from DimensionRecordStorage.insert.
135 elementRows = [record.toDict() for record in records]
136 with self._db.transaction():
137 if replace: 137 ↛ 138line 137 didn't jump to line 138, because the condition on line 137 was never true
138 self._db.replace(self._table, *elementRows)
139 elif skip_existing:
140 self._db.ensure(self._table, *elementRows, primary_key_only=True)
141 else:
142 self._db.insert(self._table, *elementRows)
143 for record in records:
144 # We really shouldn't ever get into a situation where the record
145 # here differs from the one in the DB, but the last thing we want
146 # is to make it harder to debug by making the cache different from
147 # the DB.
148 if skip_existing:
149 self._cache.setdefault(getattr(record, self.element.primaryKey.name), record)
150 else:
151 self._cache[getattr(record, self.element.primaryKey.name)] = record
152 for callback in self._callbacks:
153 callback(record)
155 def sync(self, record: DimensionRecord, update: bool = False) -> bool | dict[str, Any]:
156 # Docstring inherited from DimensionRecordStorage.sync.
157 compared = record.toDict()
158 keys = {}
159 for name in record.fields.required.names:
160 keys[name] = compared.pop(name)
161 with self._db.transaction():
162 _, inserted_or_updated = self._db.sync(
163 self._table,
164 keys=keys,
165 compared=compared,
166 update=update,
167 )
168 if inserted_or_updated: 168 ↛ 172line 168 didn't jump to line 172, because the condition on line 168 was never false
169 self._cache[getattr(record, self.element.primaryKey.name)] = record
170 for callback in self._callbacks:
171 callback(record)
172 return inserted_or_updated
174 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
175 # Docstring inherited from DimensionRecordStorage.fetch.
176 try:
177 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
178 except KeyError:
179 pass
180 # If at first we don't succeed, refresh and try again. But this time
181 # we use dict.get to return None if we don't find something.
182 self.refresh()
183 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
185 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
186 # Docstring inherited from DimensionRecordStorage.digestTables.
187 return [self._table]