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, Any, Callable, Dict, Iterable, List, Mapping, Optional
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(tsRepr=db.getTimespanRepresentation())
72 if context is not None: 72 ↛ 75line 72 didn't jump to line 75, because the condition on line 72 was never false
73 table = context.addTable(element.name, spec)
74 else:
75 table = db.ensureTableExists(element.name, spec)
76 return cls(db, element, table)
78 @property
79 def element(self) -> GovernorDimension:
80 # Docstring inherited from DimensionRecordStorage.element.
81 return self._dimension
83 def refresh(self) -> None:
84 # Docstring inherited from GovernorDimensionRecordStorage.
85 RecordClass = self._dimension.RecordClass
86 sql = sqlalchemy.sql.select(
87 [self._table.columns[name] for name in RecordClass.fields.standard.names]
88 ).select_from(self._table)
89 cache: Dict[str, DimensionRecord] = {}
90 for row in self._db.query(sql):
91 record = RecordClass(**dict(row))
92 cache[getattr(record, self._dimension.primaryKey.name)] = record
93 self._cache = cache
95 @property
96 def values(self) -> AbstractSet[str]:
97 # Docstring inherited from GovernorDimensionRecordStorage.
98 return self._cache.keys()
100 @property
101 def table(self) -> sqlalchemy.schema.Table:
102 return self._table
104 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
105 # Docstring inherited from GovernorDimensionRecordStorage.
106 self._callbacks.append(callback)
108 def clearCaches(self) -> None:
109 # Docstring inherited from DimensionRecordStorage.clearCaches.
110 self._cache.clear()
112 def join(
113 self,
114 builder: QueryBuilder, *,
115 regions: Optional[NamedKeyDict[DimensionElement, sqlalchemy.sql.ColumnElement]] = None,
116 timespans: Optional[NamedKeyDict[DimensionElement, TimespanDatabaseRepresentation]] = None,
117 ) -> None:
118 # Docstring inherited from DimensionRecordStorage.
119 joinOn = builder.startJoin(self._table, self.element.dimensions,
120 self.element.RecordClass.fields.dimensions.names)
121 builder.finishJoin(self._table, joinOn)
122 return self._table
124 def insert(self, *records: DimensionRecord) -> None:
125 # Docstring inherited from DimensionRecordStorage.insert.
126 elementRows = [record.toDict() for record in records]
127 with self._db.transaction():
128 self._db.insert(self._table, *elementRows)
129 for record in records:
130 self._cache[getattr(record, self.element.primaryKey.name)] = record
131 for callback in self._callbacks:
132 callback(record)
134 def sync(self, record: DimensionRecord) -> bool:
135 # Docstring inherited from DimensionRecordStorage.sync.
136 compared = record.toDict()
137 keys = {}
138 for name in record.fields.required.names:
139 keys[name] = compared.pop(name)
140 with self._db.transaction():
141 _, inserted = self._db.sync(
142 self._table,
143 keys=keys,
144 compared=compared,
145 )
146 if inserted:
147 self._cache[getattr(record, self.element.primaryKey.name)] = record
148 for callback in self._callbacks:
149 callback(record)
150 return inserted
152 def fetch(self, dataIds: DataCoordinateIterable) -> Iterable[DimensionRecord]:
153 # Docstring inherited from DimensionRecordStorage.fetch.
154 try:
155 return [self._cache[dataId[self.element]] for dataId in dataIds] # type: ignore
156 except KeyError:
157 pass
158 # If at first we don't succeed, refresh and try again. But this time
159 # we use dict.get to return None if we don't find something.
160 self.refresh()
161 return [self._cache.get(dataId[self.element]) for dataId in dataIds] # type: ignore
163 def digestTables(self) -> Iterable[sqlalchemy.schema.Table]:
164 # Docstring inherited from DimensionRecordStorage.digestTables.
165 return [self._table]