Coverage for python/lsst/daf/butler/registry/dimensions/governor.py: 93%
80 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-19 02:06 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-19 02:06 -0800
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, Mapping
26from typing import Any, cast
28import sqlalchemy
29from lsst.daf.relation import Relation
31from ...core import DataCoordinate, DimensionRecord, GovernorDimension
32from .. import queries
33from ..interfaces import Database, GovernorDimensionRecordStorage, StaticTablesContext
36class BasicGovernorDimensionRecordStorage(GovernorDimensionRecordStorage):
37 """A record storage implementation for `GovernorDimension` that
38 aggressively fetches and caches all values from the database.
40 Parameters
41 ----------
42 db : `Database`
43 Interface to the database engine and namespace that will hold these
44 dimension records.
45 dimension : `GovernorDimension`
46 The dimension whose records this storage will manage.
47 table : `sqlalchemy.schema.Table`
48 The logical table for the dimension.
49 """
51 def __init__(
52 self,
53 db: Database,
54 dimension: GovernorDimension,
55 table: sqlalchemy.schema.Table,
56 ):
57 self._db = db
58 self._dimension = dimension
59 self._table = table
60 # We need to allow the cache to be None so we have some recourse when
61 # it is cleared as part of transaction rollback - we can't run
62 # queries to repopulate them at that point, so we need to defer it
63 # until next use.
64 self._cache: dict[DataCoordinate, DimensionRecord] | None = None
65 self._callbacks: list[Callable[[DimensionRecord], None]] = []
67 @classmethod
68 def initialize(
69 cls,
70 db: Database,
71 element: GovernorDimension,
72 *,
73 context: StaticTablesContext | None = None,
74 config: Mapping[str, Any],
75 ) -> GovernorDimensionRecordStorage:
76 # Docstring inherited from GovernorDimensionRecordStorage.
77 spec = element.RecordClass.fields.makeTableSpec(
78 TimespanReprClass=db.getTimespanRepresentation(),
79 )
80 if context is not None: 80 ↛ 83line 80 didn't jump to line 83, because the condition on line 80 was never false
81 table = context.addTable(element.name, spec)
82 else:
83 table = db.ensureTableExists(element.name, spec)
84 return cls(db, element, table)
86 @property
87 def element(self) -> GovernorDimension:
88 # Docstring inherited from DimensionRecordStorage.element.
89 return self._dimension
91 @property
92 def table(self) -> sqlalchemy.schema.Table:
93 return self._table
95 def registerInsertionListener(self, callback: Callable[[DimensionRecord], None]) -> None:
96 # Docstring inherited from GovernorDimensionRecordStorage.
97 self._callbacks.append(callback)
99 def clearCaches(self) -> None:
100 # Docstring inherited from DimensionRecordStorage.clearCaches.
101 self._cache = None
103 def make_relation(self, context: queries.SqlQueryContext, _sized: bool = True) -> Relation:
104 # Docstring inherited.
105 payload = self._build_sql_payload(self._table, context.column_types)
106 if _sized:
107 cache = self.get_record_cache(context)
108 return context.sql_engine.make_leaf(
109 payload.columns_available.keys(),
110 name=self.element.name,
111 payload=payload,
112 min_rows=len(cache) if _sized else 0,
113 max_rows=len(cache) if _sized else None,
114 )
116 def insert(self, *records: DimensionRecord, replace: bool = False, skip_existing: bool = False) -> None:
117 # Docstring inherited from DimensionRecordStorage.insert.
118 elementRows = [record.toDict() for record in records]
119 with self._db.transaction():
120 if replace: 120 ↛ 121line 120 didn't jump to line 121, because the condition on line 120 was never true
121 self._db.replace(self._table, *elementRows)
122 elif skip_existing:
123 self._db.ensure(self._table, *elementRows, primary_key_only=True)
124 else:
125 self._db.insert(self._table, *elementRows)
126 for record in records:
127 # We really shouldn't ever get into a situation where the
128 # record here differs from the one in the DB, but the last
129 # thing we want is to make it harder to debug by making the
130 # cache different from the DB.
131 if self._cache is not None:
132 # We really shouldn't ever get into a situation where the
133 # record here differs from the one in the DB, but the last
134 # thing we want is to make it harder to debug by making the
135 # cache different from the DB.
136 if skip_existing:
137 self._cache.setdefault(record.dataId, record)
138 else:
139 self._cache[record.dataId] = record
140 for callback in self._callbacks:
141 callback(record)
143 def sync(self, record: DimensionRecord, update: bool = False) -> bool | dict[str, Any]:
144 # Docstring inherited from DimensionRecordStorage.sync.
145 compared = record.toDict()
146 keys = {}
147 for name in record.fields.required.names:
148 keys[name] = compared.pop(name)
149 with self._db.transaction():
150 _, inserted_or_updated = self._db.sync(
151 self._table,
152 keys=keys,
153 compared=compared,
154 update=update,
155 )
156 if inserted_or_updated: 156 ↛ 161line 156 didn't jump to line 161, because the condition on line 156 was never false
157 if self._cache is not None: 157 ↛ 158line 157 didn't jump to line 158, because the condition on line 157 was never true
158 self._cache[record.dataId] = record
159 for callback in self._callbacks:
160 callback(record)
161 return inserted_or_updated
163 def fetch_one(self, data_id: DataCoordinate, context: queries.SqlQueryContext) -> DimensionRecord | None:
164 # Docstring inherited.
165 cache = self.get_record_cache(context)
166 return cache.get(data_id)
168 def get_record_cache(self, context: queries.SqlQueryContext) -> Mapping[DataCoordinate, DimensionRecord]:
169 # Docstring inherited.
170 if self._cache is None:
171 reader = queries.DimensionRecordReader(self.element)
172 cache = {}
173 for row in context.fetch_iterable(self.make_relation(context, _sized=False)):
174 record = reader.read(row)
175 cache[record.dataId] = record
176 self._cache = cache
177 return cast(Mapping[DataCoordinate, DimensionRecord], self._cache)
179 def digestTables(self) -> list[sqlalchemy.schema.Table]:
180 # Docstring inherited from DimensionRecordStorage.digestTables.
181 return [self._table]