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

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, Callable, Dict, Iterable, List, Optional
27import sqlalchemy
29from ...core import (
30 Config,
31 DataCoordinateIterable,
32 DimensionElement,
33 DimensionRecord,
34 GovernorDimension,
35 NamedKeyDict,
36 TimespanDatabaseRepresentation,
37)
38from ..interfaces import (
39 Database,
40 GovernorDimensionRecordStorage,
41 StaticTablesContext,
42)
43from ..queries import QueryBuilder
46class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage):
47 """A record storage implementation for `GovernorDimension` that
48 aggressively fetches and caches all values from the database.
50 Parameters
51 ----------
52 db : `Database`
53 Interface to the database engine and namespace that will hold these
54 dimension records.
55 dimension : `GovernorDimension`
56 The dimension whose records this storage will manage.
57 table : `sqlalchemy.schema.Table`
58 The logical table for the dimension.
59 """
60 def __init__(self, db: Database, dimension: GovernorDimension, table: sqlalchemy.schema.Table):
61 self._db = db
62 self._dimension = dimension
63 self._table = table
64 self._cache: Dict[str, DimensionRecord] = {}
65 self._callbacks: List[Callable[[DimensionRecord], None]] = []
67 @classmethod
68 def initialize(cls, db: Database, element: GovernorDimension, *,
69 context: Optional[StaticTablesContext] = None,
70 config: Config) -> GovernorDimensionRecordStorage:
71 # Docstring inherited from GovernorDimensionRecordStorage.
72 spec = element.RecordClass.fields.makeTableSpec(tsRepr=db.getTimespanRepresentation())
73 if context is not None: 73 ↛ 76line 73 didn't jump to line 76, because the condition on line 73 was never false
74 table = context.addTable(element.name, spec)
75 else:
76 table = db.ensureTableExists(element.name, spec)
77 return cls(db, element, table)
79 @property
80 def element(self) -> GovernorDimension:
81 # Docstring inherited from DimensionRecordStorage.element.
82 return self._dimension
84 def refresh(self) -> None:
85 # Docstring inherited from GovernorDimensionRecordStorage.
86 RecordClass = self._dimension.RecordClass
87 sql = sqlalchemy.sql.select(
88 [self._table.columns[name] for name in RecordClass.fields.standard.names]
89 ).select_from(self._table)
90 cache: Dict[str, DimensionRecord] = {}
91 for row in self._db.query(sql):
92 record = RecordClass(**dict(row))
93 cache[getattr(record, self._dimension.primaryKey.name)] = record
94 self._cache = cache
96 @property
97 def values(self) -> AbstractSet[str]:
98 # Docstring inherited from GovernorDimensionRecordStorage.
99 return self._cache.keys()
101 @property
102 def table(self) -> sqlalchemy.schema.Table:
103 return self._table
105 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
106 # Docstring inherited from GovernorDimensionRecordStorage.
107 self._callbacks.append(callback)
109 def clearCaches(self) -> None:
110 # Docstring inherited from DimensionRecordStorage.clearCaches.
111 self._cache.clear()
113 def join(
114 self,
115 builder: QueryBuilder, *,
116 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
117 timespans: Optional[NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation]] = None,
118 ) -> None:
119 # Docstring inherited from DimensionRecordStorage.
120 joinOn = builder.startJoin(self._table, self.element.dimensions,
121 self.element.RecordClass.fields.dimensions.names)
122 builder.finishJoin(self._table, joinOn)
123 return self._table
125 def insert(self, *records: DimensionRecord) -> None:
126 # Docstring inherited from DimensionRecordStorage.insert.
127 elementRows = [record.toDict() for record in records]
128 with self._db.transaction():
129 self._db.insert(self._table, *elementRows)
130 for record in records:
131 self._cache[getattr(record, self.element.primaryKey.name)] = record
132 for callback in self._callbacks:
133 callback(record)
135 def sync(self, record: DimensionRecord) -> bool:
136 # Docstring inherited from DimensionRecordStorage.sync.
137 compared = record.toDict()
138 keys = {}
139 for name in record.fields.required.names:
140 keys[name] = compared.pop(name)
141 with self._db.transaction():
142 _, inserted = self._db.sync(
143 self._table,
144 keys=keys,
145 compared=compared,
146 )
147 if inserted:
148 self._cache[getattr(record, self.element.primaryKey.name)] = record
149 for callback in self._callbacks:
150 callback(record)
151 return inserted
153 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
154 # Docstring inherited from DimensionRecordStorage.fetch.
155 try:
156 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
157 except KeyError:
158 pass
159 # If at first we don't succeed, refresh and try again. But this time
160 # we use dict.get to return None if we don't find something.
161 self.refresh()
162 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
164 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
165 # Docstring inherited from DimensionRecordStorage.digestTables.
166 return [self._table]