Coverage for python/lsst/daf/butler/tests/_dummyRegistry.py: 28%
107 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-10-25 15:14 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-10-25 15:14 +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 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__ = ("DummyRegistry",)
25from collections.abc import Iterable, Iterator
26from typing import Any
28import sqlalchemy
29from lsst.daf.butler import DimensionUniverse, ddl
30from lsst.daf.butler.registry.bridge.ephemeral import EphemeralDatastoreRegistryBridge
31from lsst.daf.butler.registry.interfaces import (
32 Database,
33 DatabaseInsertMode,
34 DatasetIdRef,
35 DatasetRecordStorageManager,
36 DatastoreRegistryBridge,
37 DatastoreRegistryBridgeManager,
38 OpaqueTableStorage,
39 OpaqueTableStorageManager,
40 StaticTablesContext,
41 VersionTuple,
42)
44from ..core.datastore import DatastoreTransaction
47class DummyOpaqueTableStorage(OpaqueTableStorage):
48 def __init__(self, name: str, spec: ddl.TableSpec) -> None:
49 super().__init__(name=name)
50 self._rows: list[dict] = []
51 self._spec = spec
53 def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
54 # Docstring inherited from OpaqueTableStorage.
55 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.INSERT)
57 def replace(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
58 # Docstring inherited from OpaqueTableStorage.
59 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.REPLACE)
61 def ensure(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
62 # Docstring inherited from OpaqueTableStorage.
63 self._insert(*data, transaction=transaction, insert_mode=DatabaseInsertMode.ENSURE)
65 def _insert(
66 self,
67 *data: dict,
68 transaction: DatastoreTransaction | None = None,
69 insert_mode: DatabaseInsertMode = DatabaseInsertMode.INSERT,
70 ) -> None:
71 uniqueConstraints = list(self._spec.unique)
72 uniqueConstraints.append(tuple(field.name for field in self._spec.fields if field.primaryKey))
73 for d in data:
74 skipping = False
75 for constraint in uniqueConstraints:
76 matching = list(self.fetch(**{k: d[k] for k in constraint}))
77 if len(matching) != 0:
78 match insert_mode:
79 case DatabaseInsertMode.INSERT:
80 raise RuntimeError(
81 f"Unique constraint {constraint} violation in external table {self.name}."
82 )
83 case DatabaseInsertMode.ENSURE:
84 # Row already exists. Skip.
85 skipping = True
86 case DatabaseInsertMode.REPLACE:
87 # Should try to put these rows back on transaction
88 # rollback...
89 self.delete([], *matching)
90 case _:
91 raise ValueError(f"Unrecognized insert mode: {insert_mode}.")
93 if skipping:
94 continue
95 self._rows.append(d)
96 if transaction is not None:
97 transaction.registerUndo("insert", self.delete, [], d)
99 def fetch(self, **where: Any) -> Iterator[dict]:
100 # Docstring inherited from OpaqueTableStorage.
101 where = where.copy() # May need to modify it.
103 # Can support an IN operator if given list.
104 wherein = {}
105 for k in list(where):
106 if isinstance(where[k], tuple | list | set):
107 wherein[k] = set(where[k])
108 del where[k]
110 for d in self._rows:
111 if all(d[k] == v for k, v in where.items()):
112 if wherein:
113 match = True
114 for k, v in wherein.items():
115 if d[k] not in v:
116 match = False
117 break
118 if match:
119 yield d
120 else:
121 yield d
123 def delete(self, columns: Iterable[str], *rows: dict) -> None:
124 # Docstring inherited from OpaqueTableStorage.
125 kept_rows = []
126 for table_row in self._rows:
127 for where_row in rows:
128 if all(table_row[k] == v for k, v in where_row.items()):
129 break
130 else:
131 kept_rows.append(table_row)
132 self._rows = kept_rows
135class DummyOpaqueTableStorageManager(OpaqueTableStorageManager):
136 def __init__(self, registry_schema_version: VersionTuple | None = None) -> None:
137 super().__init__(registry_schema_version=registry_schema_version)
138 self._storages: dict[str, DummyOpaqueTableStorage] = {}
140 @classmethod
141 def initialize(
142 cls, db: Database, context: StaticTablesContext, registry_schema_version: VersionTuple | None = None
143 ) -> OpaqueTableStorageManager:
144 # Docstring inherited from OpaqueTableStorageManager.
145 # Not used, but needed to satisfy ABC requirement.
146 return cls(registry_schema_version=registry_schema_version)
148 def get(self, name: str) -> OpaqueTableStorage | None:
149 # Docstring inherited from OpaqueTableStorageManager.
150 return self._storages.get(name)
152 def register(self, name: str, spec: ddl.TableSpec) -> OpaqueTableStorage:
153 # Docstring inherited from OpaqueTableStorageManager.
154 return self._storages.setdefault(name, DummyOpaqueTableStorage(name, spec))
156 @classmethod
157 def currentVersions(cls) -> list[VersionTuple]:
158 # Docstring inherited from VersionedExtension.
159 return []
162class DummyDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager):
163 def __init__(
164 self,
165 opaque: OpaqueTableStorageManager,
166 universe: DimensionUniverse,
167 datasetIdColumnType: type,
168 registry_schema_version: VersionTuple | None = None,
169 ):
170 super().__init__(
171 opaque=opaque,
172 universe=universe,
173 datasetIdColumnType=datasetIdColumnType,
174 registry_schema_version=registry_schema_version,
175 )
176 self._bridges: dict[str, EphemeralDatastoreRegistryBridge] = {}
178 @classmethod
179 def initialize(
180 cls,
181 db: Database,
182 context: StaticTablesContext,
183 *,
184 opaque: OpaqueTableStorageManager,
185 datasets: type[DatasetRecordStorageManager],
186 universe: DimensionUniverse,
187 registry_schema_version: VersionTuple | None = None,
188 ) -> DatastoreRegistryBridgeManager:
189 # Docstring inherited from DatastoreRegistryBridgeManager
190 # Not used, but needed to satisfy ABC requirement.
191 return cls(
192 opaque=opaque,
193 universe=universe,
194 datasetIdColumnType=datasets.getIdColumnType(),
195 registry_schema_version=registry_schema_version,
196 )
198 def refresh(self) -> None:
199 # Docstring inherited from DatastoreRegistryBridgeManager
200 pass
202 def register(self, name: str, *, ephemeral: bool = False) -> DatastoreRegistryBridge:
203 # Docstring inherited from DatastoreRegistryBridgeManager
204 return self._bridges.setdefault(name, EphemeralDatastoreRegistryBridge(name))
206 def findDatastores(self, ref: DatasetIdRef) -> Iterable[str]:
207 # Docstring inherited from DatastoreRegistryBridgeManager
208 for name, bridge in self._bridges.items():
209 if ref in bridge:
210 yield name
212 @classmethod
213 def currentVersions(cls) -> list[VersionTuple]:
214 # Docstring inherited from VersionedExtension.
215 return []
218class DummyRegistry:
219 """Dummy Registry, for Datastore test purposes."""
221 def __init__(self) -> None:
222 self._opaque = DummyOpaqueTableStorageManager()
223 self.dimensions = DimensionUniverse()
224 self._datastoreBridges = DummyDatastoreRegistryBridgeManager(
225 self._opaque, self.dimensions, sqlalchemy.BigInteger
226 )
228 def getDatastoreBridgeManager(self) -> DatastoreRegistryBridgeManager:
229 return self._datastoreBridges