Coverage for python/lsst/daf/butler/registry/dimensions/governor.py: 95%
82 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-07 09:46 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2022-10-07 09:46 +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 for row in self._db.query(sql):
96 record = RecordClass(**row._asdict())
97 cache[getattr(record, self._dimension.primaryKey.name)] = record
98 self._cache = cache
100 @property
101 def values(self) -> Set[str]:
102 # Docstring inherited from GovernorDimensionRecordStorage.
103 return self._cache.keys()
105 @property
106 def table(self) -> sqlalchemy.schema.Table:
107 return self._table
109 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
110 # Docstring inherited from GovernorDimensionRecordStorage.
111 self._callbacks.append(callback)
113 def clearCaches(self) -> None:
114 # Docstring inherited from DimensionRecordStorage.clearCaches.
115 self._cache.clear()
117 def join(
118 self,
119 builder: QueryBuilder,
120 *,
121 regions: NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement] | None = None,
122 timespans: NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation] | None = None,
123 ) -> None:
124 # Docstring inherited from DimensionRecordStorage.
125 joinOn = builder.startJoin(
126 self._table, self.element.dimensions, self.element.RecordClass.fields.dimensions.names
127 )
128 builder.finishJoin(self._table, joinOn)
129 return self._table
131 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None:
132 # Docstring inherited from DimensionRecordStorage.insert.
133 elementRows = [record.toDict() for record in records]
134 with self._db.transaction():
135 if replace: 135 ↛ 136line 135 didn't jump to line 136, because the condition on line 135 was never true
136 self._db.replace(self._table, *elementRows)
137 elif skip_existing:
138 self._db.ensure(self._table, *elementRows, primary_key_only=True)
139 else:
140 self._db.insert(self._table, *elementRows)
141 for record in records:
142 # We really shouldn't ever get into a situation where the record
143 # here differs from the one in the DB, but the last thing we want
144 # is to make it harder to debug by making the cache different from
145 # the DB.
146 if skip_existing:
147 self._cache.setdefault(getattr(record, self.element.primaryKey.name), record)
148 else:
149 self._cache[getattr(record, self.element.primaryKey.name)] = record
150 for callback in self._callbacks:
151 callback(record)
153 def sync(self, record: DimensionRecord, update: bool = False) -> bool | dict[str, Any]:
154 # Docstring inherited from DimensionRecordStorage.sync.
155 compared = record.toDict()
156 keys = {}
157 for name in record.fields.required.names:
158 keys[name] = compared.pop(name)
159 with self._db.transaction():
160 _, inserted_or_updated = self._db.sync(
161 self._table,
162 keys=keys,
163 compared=compared,
164 update=update,
165 )
166 if inserted_or_updated: 166 ↛ 170line 166 didn't jump to line 170, because the condition on line 166 was never false
167 self._cache[getattr(record, self.element.primaryKey.name)] = record
168 for callback in self._callbacks:
169 callback(record)
170 return inserted_or_updated
172 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
173 # Docstring inherited from DimensionRecordStorage.fetch.
174 try:
175 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
176 except KeyError:
177 pass
178 # If at first we don't succeed, refresh and try again. But this time
179 # we use dict.get to return None if we don't find something.
180 self.refresh()
181 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
183 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
184 # Docstring inherited from DimensionRecordStorage.digestTables.
185 return [self._table]