Coverage for python/lsst/daf/butler/tests/_dummyRegistry.py: 28%
107 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-06 10:53 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-12-06 10:53 +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 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 @classmethod
147 def initialize(
148 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None
149 ) -> OpaqueTableStorageManager:
150 # Docstring inherited from OpaqueTableStorageManager.
151 # Not used, but needed to satisfy ABC requirement.
152 return cls(registry_schema_version=registry_schema_version)
154 def get(self, name: str) -> OpaqueTableStorage | None:
155 # Docstring inherited from OpaqueTableStorageManager.
156 return self._storages.get(name)
158 def register(self, name: str, spec: ddl.TableSpec) -> OpaqueTableStorage:
159 # Docstring inherited from OpaqueTableStorageManager.
160 return self._storages.setdefault(name, DummyOpaqueTableStorage(name, spec))
162 @classmethod
163 def currentVersions(cls) -> list[VersionTuple]:
164 # Docstring inherited from VersionedExtension.
165 return []
168class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager):
169 def __init__(
170 self,
171 opaque: OpaqueTableStorageManager,
172 universe: DimensionUniverse,
173 datasetIdColumnType: type,
174 registry_schema_version: VersionTuple | None = None,
175 ):
176 super().__init__(
177 opaque=opaque,
178 universe=universe,
179 datasetIdColumnType=datasetIdColumnType,
180 registry_schema_version=registry_schema_version,
181 )
182 self._bridges: dict[str, EphemeralDatastoreRegistryBridge] = {}
184 @classmethod
185 def initialize(
186 cls,
187 db: Database,
188 context: StaticTablesContext,
189 *,
190 opaque: OpaqueTableStorageManager,
191 datasets: type[DatasetRecordStorageManager],
192 universe: DimensionUniverse,
193 registry_schema_version: VersionTuple | None = None,
194 ) -> DatastoreRegistryBridgeManager:
195 # Docstring inherited from DatastoreRegistryBridgeManager
196 # Not used, but needed to satisfy ABC requirement.
197 return cls(
198 opaque=opaque,
199 universe=universe,
200 datasetIdColumnType=datasets.getIdColumnType(),
201 registry_schema_version=registry_schema_version,
202 )
204 def refresh(self) -> None:
205 # Docstring inherited from DatastoreRegistryBridgeManager
206 pass
208 def register(self, name: str, *, ephemeral: bool = False) -> DatastoreRegistryBridge:
209 # Docstring inherited from DatastoreRegistryBridgeManager
210 return self._bridges.setdefault(name, EphemeralDatastoreRegistryBridge(name))
212 def findDatastores(self, ref: DatasetIdRef) -> Iterable[str]:
213 # Docstring inherited from DatastoreRegistryBridgeManager
214 for name, bridge in self._bridges.items():
215 if ref in bridge:
216 yield name
218 @classmethod
219 def currentVersions(cls) -> list[VersionTuple]:
220 # Docstring inherited from VersionedExtension.
221 return []
224class DummyRegistry:
225 """Dummy Registry, for Datastore test purposes."""
227 def __init__(self) -> None:
228 self._opaque = DummyOpaqueTableStorageManager()
229 self.dimensions = DimensionUniverse()
230 self._datastoreBridges = DummyDatastoreRegistryBridgeManager(
231 self._opaque, self.dimensions, sqlalchemy.BigInteger
232 )
234 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager:
235 return self._datastoreBridges