Coverage for python/lsst/daf/butler/tests/_dummyRegistry.py: 28%
111 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:53 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:53 -0700
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = ("DummyRegistry",)
31from collections.abc import Iterable, Iterator
32from typing import Any
34import sqlalchemy
35from lsst.daf.butler import DimensionUniverse, ddl
36from lsst.daf.butler.registry.bridge.ephemeral import EphemeralDatastoreRegistryBridge
37from lsst.daf.butler.registry.interfaces import (
38 Database,
39 DatabaseInsertMode,
40 DatasetIdRef,
41 DatasetRecordStorageManager,
42 DatastoreRegistryBridge,
43 DatastoreRegistryBridgeManager,
44 OpaqueTableStorage,
45 OpaqueTableStorageManager,
46 StaticTablesContext,
47 VersionTuple,
48)
50from ..datastore import DatastoreTransaction
53class DummyOpaqueTableStorage(OpaqueTableStorage):
54 def __init__(self, name: str, spec: ddl.TableSpec) -> None:
55 super().__init__(name=name)
56 self._rows: list[dict] = []
57 self._spec = spec
59 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
60 # Docstring inherited from OpaqueTableStorage.
61 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.INSERT)
63 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
64 # Docstring inherited from OpaqueTableStorage.
65 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.REPLACE)
67 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
68 # Docstring inherited from OpaqueTableStorage.
69 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.ENSURE)
71 def _insert(
72 self,
73 *data: dict,
74 transaction: DatastoreTransaction | None = None,
75 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT,
76 ) -> None:
77 uniqueConstraints = list(self._spec.unique)
78 uniqueConstraints.append(tuple(field.name for field in self._spec.fields if field.primaryKey))
79 for d in data:
80 skipping = False
81 for constraint in uniqueConstraints:
82 matching = list(self.fetch(**{k: d[k] for k in constraint}))
83 if len(matching) != 0:
84 match insert_mode:
85 case DatabaseInsertMode.INSERT:
86 raise RuntimeError(
87 f"Unique constraint {constraint} violation in external table {self.name}."
88 )
89 case DatabaseInsertMode.ENSURE:
90 # Row already exists. Skip.
91 skipping = True
92 case DatabaseInsertMode.REPLACE:
93 # Should try to put these rows back on transaction
94 # rollback...
95 self.delete([], *matching)
96 case _:
97 raise ValueError(f"Unrecognized insert mode: {insert_mode}.")
99 if skipping:
100 continue
101 self._rows.append(d)
102 if transaction is not None:
103 transaction.registerUndo("insert", self.delete, [], d)
105 def fetch(self, **where: Any) -> Iterator[dict]:
106 # Docstring inherited from OpaqueTableStorage.
107 where = where.copy() # May need to modify it.
109 # Can support an IN operator if given list.
110 wherein = {}
111 for k in list(where):
112 if isinstance(where[k], tuple | list | set):
113 wherein[k] = set(where[k])
114 del where[k]
116 for d in self._rows:
117 if all(d[k] == v for k, v in where.items()):
118 if wherein:
119 match = True
120 for k, v in wherein.items():
121 if d[k] not in v:
122 match = False
123 break
124 if match:
125 yield d
126 else:
127 yield d
129 def delete(self, columns: Iterable[str], *rows: dict) -> None:
130 # Docstring inherited from OpaqueTableStorage.
131 kept_rows = []
132 for table_row in self._rows:
133 for where_row in rows:
134 if all(table_row[k] == v for k, v in where_row.items()):
135 break
136 else:
137 kept_rows.append(table_row)
138 self._rows = kept_rows
141class DummyOpaqueTableStorageManager(OpaqueTableStorageManager):
142 def __init__(self, registry_schema_version: VersionTuple | None = None) -> None:
143 super().__init__(registry_schema_version=registry_schema_version)
144 self._storages: dict[str, DummyOpaqueTableStorage] = {}
146 def clone(self, db: Database) -> OpaqueTableStorageManager:
147 return self
149 @classmethod
150 def initialize(
151 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None
152 ) -> OpaqueTableStorageManager:
153 # Docstring inherited from OpaqueTableStorageManager.
154 # Not used, but needed to satisfy ABC requirement.
155 return cls(registry_schema_version=registry_schema_version)
157 def get(self, name: str) -> OpaqueTableStorage | None:
158 # Docstring inherited from OpaqueTableStorageManager.
159 return self._storages.get(name)
161 def register(self, name: str, spec: ddl.TableSpec) -> OpaqueTableStorage:
162 # Docstring inherited from OpaqueTableStorageManager.
163 return self._storages.setdefault(name, DummyOpaqueTableStorage(name, spec))
165 @classmethod
166 def currentVersions(cls) -> list[VersionTuple]:
167 # Docstring inherited from VersionedExtension.
168 return []
171class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager):
172 def __init__(
173 self,
174 opaque: OpaqueTableStorageManager,
175 universe: DimensionUniverse,
176 datasetIdColumnType: type,
177 registry_schema_version: VersionTuple | None = None,
178 ):
179 super().__init__(
180 opaque=opaque,
181 universe=universe,
182 datasetIdColumnType=datasetIdColumnType,
183 registry_schema_version=registry_schema_version,
184 )
185 self._bridges: dict[str, EphemeralDatastoreRegistryBridge] = {}
187 def clone(self, *, db: Database, opaque: OpaqueTableStorageManager) -> DatastoreRegistryBridgeManager:
188 return DummyDatastoreRegistryBridgeManager(
189 opaque=opaque,
190 universe=self.universe,
191 datasetIdColumnType=self.datasetIdColumnType,
192 registry_schema_version=self._registry_schema_version,
193 )
195 @classmethod
196 def initialize(
197 cls,
198 db: Database,
199 context: StaticTablesContext,
200 *,
201 opaque: OpaqueTableStorageManager,
202 datasets: type[DatasetRecordStorageManager],
203 universe: DimensionUniverse,
204 registry_schema_version: VersionTuple | None = None,
205 ) -> DatastoreRegistryBridgeManager:
206 # Docstring inherited from DatastoreRegistryBridgeManager
207 # Not used, but needed to satisfy ABC requirement.
208 return cls(
209 opaque=opaque,
210 universe=universe,
211 datasetIdColumnType=datasets.getIdColumnType(),
212 registry_schema_version=registry_schema_version,
213 )
215 def refresh(self) -> None:
216 # Docstring inherited from DatastoreRegistryBridgeManager
217 pass
219 def register(self, name: str, *, ephemeral: bool = False) -> DatastoreRegistryBridge:
220 # Docstring inherited from DatastoreRegistryBridgeManager
221 return self._bridges.setdefault(name, EphemeralDatastoreRegistryBridge(name))
223 def findDatastores(self, ref: DatasetIdRef) -> Iterable[str]:
224 # Docstring inherited from DatastoreRegistryBridgeManager
225 for name, bridge in self._bridges.items():
226 if ref in bridge:
227 yield name
229 @classmethod
230 def currentVersions(cls) -> list[VersionTuple]:
231 # Docstring inherited from VersionedExtension.
232 return []
235class DummyRegistry:
236 """Dummy Registry, for Datastore test purposes."""
238 def __init__(self) -> None:
239 self._opaque = DummyOpaqueTableStorageManager()
240 self.dimensions = DimensionUniverse()
241 self._datastoreBridges = DummyDatastoreRegistryBridgeManager(
242 self._opaque, self.dimensions, sqlalchemy.BigInteger
243 )
245 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager:
246 return self._datastoreBridges