Coverage for python/lsst/daf/butler/registry/bridge/monolithic.py: 84%
126 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-06 01:43 -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 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
29from ... import ddl
31__all__ = ("MonolithicDatastoreRegistryBridge", "MonolithicDatastoreRegistryBridgeManager")
33from collections import namedtuple
34from collections.abc import Collection, Iterable, Iterator
35from contextlib import contextmanager
36from typing import TYPE_CHECKING, cast
38import sqlalchemy
40from lsst.utils.iteration import chunk_iterable
42from ..._dataset_ref import DatasetId
43from ...datastore.stored_file_info import StoredDatastoreItemInfo
44from .._exceptions import ConflictingDefinitionError
45from ..interfaces import (
46 DatasetIdRef,
47 DatastoreRegistryBridge,
48 DatastoreRegistryBridgeManager,
49 FakeDatasetRef,
50 OpaqueTableStorage,
51 VersionTuple,
52)
53from ..opaque import ByNameOpaqueTableStorage
54from .ephemeral import EphemeralDatastoreRegistryBridge
56if TYPE_CHECKING:
57 from ...datastore import DatastoreTransaction
58 from ...dimensions import DimensionUniverse
59 from ..interfaces import (
60 Database,
61 DatasetRecordStorageManager,
62 OpaqueTableStorageManager,
63 StaticTablesContext,
64 )
66_TablesTuple = namedtuple(
67 "_TablesTuple",
68 [
69 "dataset_location",
70 "dataset_location_trash",
71 ],
72)
74# This has to be updated on every schema change
75_VERSION = VersionTuple(0, 2, 1)
78def _makeTableSpecs(datasets: type[DatasetRecordStorageManager]) -> _TablesTuple:
79 """Construct specifications for tables used by the monolithic datastore
80 bridge classes.
82 Parameters
83 ----------
84 datasets : subclass of `DatasetRecordStorageManager`
85 Manager class for datasets; used only to create foreign key fields.
87 Returns
88 -------
89 specs : `_TablesTuple`
90 A named tuple containing `ddl.TableSpec` instances.
91 """
92 # We want the dataset_location and dataset_location_trash tables
93 # to have the same definition, aside from the behavior of their link
94 # to the dataset table: the trash table has no foreign key constraint.
95 # The order of columns in dataset_location_trash is reversed, it is more
96 # optimal for query planner.
98 datastore_field = ddl.FieldSpec(
99 name="datastore_name",
100 dtype=sqlalchemy.String,
101 length=256,
102 primaryKey=True,
103 nullable=False,
104 doc="Name of the Datastore this entry corresponds to.",
105 )
107 dataset_location = ddl.TableSpec(
108 doc=(
109 "A table that provides information on whether a dataset is stored in "
110 "one or more Datastores. The presence or absence of a record in this "
111 "table itself indicates whether the dataset is present in that "
112 "Datastore. "
113 ),
114 fields=[datastore_field],
115 )
116 datasets.addDatasetForeignKey(dataset_location, primaryKey=True)
118 dataset_location_trash = ddl.TableSpec(
119 doc="A table that keeps iinformation about datasets that are removed from Datastores.",
120 fields=[],
121 )
122 datasets.addDatasetForeignKey(dataset_location_trash, primaryKey=True, constraint=False)
123 dataset_location_trash.fields.add(datastore_field)
125 return _TablesTuple(
126 dataset_location=dataset_location,
127 dataset_location_trash=dataset_location_trash,
128 )
131class MonolithicDatastoreRegistryBridge(DatastoreRegistryBridge):
132 """An implementation of `DatastoreRegistryBridge` that uses the same two
133 tables for all non-ephemeral datastores.
135 Parameters
136 ----------
137 datastoreName : `str`
138 Name of the `Datastore` as it should appear in `Registry` tables
139 referencing it.
140 db : `Database`
141 Object providing a database connection and generic distractions.
142 tables : `_TablesTuple`
143 Named tuple containing `sqlalchemy.schema.Table` instances.
144 """
146 def __init__(self, datastoreName: str, *, db: Database, tables: _TablesTuple):
147 super().__init__(datastoreName)
148 self._db = db
149 self._tables = tables
151 def _refsToRows(self, refs: Iterable[DatasetIdRef]) -> list[dict]:
152 """Transform an iterable of `DatasetRef` or `FakeDatasetRef` objects to
153 a list of dictionaries that match the schema of the tables used by this
154 class.
156 Parameters
157 ----------
158 refs : `~collections.abc.Iterable` [ `DatasetRef` or `FakeDatasetRef` ]
159 Datasets to transform.
161 Returns
162 -------
163 rows : `list` [ `dict` ]
164 List of dictionaries, with "datastoreName" and "dataset_id" keys.
165 """
166 return [{"datastore_name": self.datastoreName, "dataset_id": ref.id} for ref in refs]
168 def ensure(self, refs: Iterable[DatasetIdRef]) -> None:
169 # Docstring inherited from DatastoreRegistryBridge
170 self._db.ensure(self._tables.dataset_location, *self._refsToRows(refs))
172 def insert(self, refs: Iterable[DatasetIdRef]) -> None:
173 # Docstring inherited from DatastoreRegistryBridge
174 try:
175 self._db.insert(self._tables.dataset_location, *self._refsToRows(refs))
176 except sqlalchemy.exc.IntegrityError as exc:
177 raise ConflictingDefinitionError("Some datasets were already registered.") from exc
179 def forget(self, refs: Iterable[DatasetIdRef]) -> None:
180 # Docstring inherited from DatastoreRegistryBridge
181 with self._db.transaction():
182 # The list of IDs can be very large, split it into reasonable size
183 # chunks to avoid hitting limits.
184 for refs_chunk in chunk_iterable(refs, 50_000):
185 dataset_ids = [ref.id for ref in refs_chunk]
186 where = sqlalchemy.sql.and_(
187 self._tables.dataset_location.columns.datastore_name == self.datastoreName,
188 self._db.make_in_array_constraint(
189 self._tables.dataset_location.columns.dataset_id, dataset_ids
190 ),
191 )
192 self._db.deleteWhere(self._tables.dataset_location, where)
194 def moveToTrash(self, refs: Iterable[DatasetIdRef], transaction: DatastoreTransaction | None) -> None:
195 # Docstring inherited from DatastoreRegistryBridge
196 location = self._tables.dataset_location
197 location_trash = self._tables.dataset_location_trash
198 with self._db.transaction():
199 for refs_chunk in chunk_iterable(refs, 50_000):
200 # We only want to move IDs that actually exist in the
201 # dataset_location table. Instead of querying for existing IDs,
202 # which would need an extra query, we use INSERT ... SELECT
203 # and DELETE using WHERE clause that limits operations to
204 # existing IDs.
205 dataset_ids = [ref.id for ref in refs_chunk]
207 where = sqlalchemy.sql.and_(
208 location.columns.datastore_name == self.datastoreName,
209 self._db.make_in_array_constraint(location.columns.dataset_id, dataset_ids),
210 )
212 select = (
213 sqlalchemy.sql.select(location.columns.datastore_name, location.columns.dataset_id)
214 .where(where)
215 .with_for_update()
216 )
217 self._db.insert(location_trash, select=select)
219 self._db.deleteWhere(location, where)
221 def check(self, datasets: Iterable[DatasetId]) -> set[DatasetId]:
222 # Docstring inherited from DatastoreRegistryBridge
223 found: set[DatasetId] = set()
224 with self._db.session():
225 for batch in chunk_iterable(datasets, 50000):
226 sql = (
227 sqlalchemy.sql.select(self._tables.dataset_location.columns.dataset_id)
228 .select_from(self._tables.dataset_location)
229 .where(
230 sqlalchemy.sql.and_(
231 self._tables.dataset_location.columns.datastore_name == self.datastoreName,
232 self._db.make_in_array_constraint(
233 self._tables.dataset_location.columns.dataset_id, batch
234 ),
235 )
236 )
237 )
238 with self._db.query(sql) as sql_result:
239 sql_ids = sql_result.scalars().all()
240 found.update(sql_ids)
242 return found
244 @contextmanager
245 def emptyTrash(
246 self,
247 records_table: OpaqueTableStorage | None = None,
248 record_class: type[StoredDatastoreItemInfo] | None = None,
249 record_column: str | None = None,
250 selected_ids: Collection[DatasetId] | None = None,
251 dry_run: bool = False,
252 ) -> Iterator[tuple[Iterable[tuple[DatasetIdRef, StoredDatastoreItemInfo | None]], set[str] | None]]:
253 # Docstring inherited from DatastoreRegistryBridge
255 if records_table is None: 255 ↛ 256line 255 didn't jump to line 256 because the condition on line 255 was never true
256 raise ValueError("This implementation requires a records table.")
258 assert isinstance(records_table, ByNameOpaqueTableStorage), (
259 f"Records table must support hidden attributes. Got {type(records_table)}."
260 )
262 if record_class is None: 262 ↛ 263line 262 didn't jump to line 263 because the condition on line 262 was never true
263 raise ValueError("Record class must be provided if records table is given.")
265 # Helper closure to generate the common join+where clause.
266 def join_records(
267 select: sqlalchemy.sql.Select, location_table: sqlalchemy.schema.Table
268 ) -> sqlalchemy.sql.Select:
269 # mypy needs to be sure
270 assert isinstance(records_table, ByNameOpaqueTableStorage)
271 return select.select_from(
272 records_table._table.join(
273 location_table,
274 onclause=records_table._table.columns.dataset_id == location_table.columns.dataset_id,
275 )
276 ).where(location_table.columns.datastore_name == self.datastoreName)
278 # SELECT records.dataset_id, records.path FROM records
279 # JOIN records on dataset_location.dataset_id == records.dataset_id
280 # WHERE dataset_location.datastore_name = datastoreName
282 # It's possible that we may end up with a ref listed in the trash
283 # table that is not listed in the records table. Such an
284 # inconsistency would be missed by this query.
285 info_in_trash = join_records(records_table._table.select(), self._tables.dataset_location_trash)
286 if selected_ids:
287 info_in_trash = info_in_trash.where(
288 self._db.make_in_array_constraint(
289 self._tables.dataset_location_trash.columns["dataset_id"], selected_ids
290 )
291 )
292 info_in_trash = info_in_trash.with_for_update(skip_locked=True)
294 # Run query, transform results into a list of dicts that we can later
295 # use to delete.
296 with self._db.query(info_in_trash) as sql_result:
297 rows = [dict(row, datastore_name=self.datastoreName) for row in sql_result.mappings()]
299 # It is possible for trashed refs to be linked to artifacts that
300 # are still associated with refs that are not to be trashed. We
301 # need to be careful to consider those and indicate to the caller
302 # that those artifacts should be retained. Can only do this check
303 # if the caller provides a column name that can map to multiple
304 # refs.
305 preserved: set[str] | None = None
306 if record_column is not None: 306 ↛ 344line 306 didn't jump to line 344 because the condition on line 306 was always true
307 # The artifacts belonging to the trashed refs under consideration.
308 items_in_trash = join_records(
309 sqlalchemy.sql.select(records_table._table.columns[record_column]),
310 self._tables.dataset_location_trash,
311 )
312 if selected_ids:
313 items_in_trash = items_in_trash.where(
314 self._db.make_in_array_constraint(
315 self._tables.dataset_location_trash.columns["dataset_id"], selected_ids
316 )
317 )
318 items_in_trash_alias = items_in_trash.alias("items_in_trash")
320 # Whether one of those artifacts is also referenced by a ref that
321 # is not in the trash. Written as a correlated EXISTS so that the
322 # query is driven by the trashed artifacts and probes the records
323 # table by artifact, rather than materializing every artifact the
324 # datastore holds in order to join the two sets.
325 referenced_by_live_ref = join_records(
326 sqlalchemy.sql.select(sqlalchemy.sql.literal(1)), self._tables.dataset_location
327 ).where(
328 records_table._table.columns[record_column] == items_in_trash_alias.columns[record_column]
329 )
331 # A query for paths that are referenced by datasets in the trash
332 # and datasets not in the trash.
333 items_to_preserve = (
334 sqlalchemy.sql.select(items_in_trash_alias.columns[record_column])
335 .distinct()
336 .where(sqlalchemy.exists(referenced_by_live_ref))
337 )
338 with self._db.query(items_to_preserve) as sql_result:
339 preserved = {row[record_column] for row in sql_result.mappings()}
341 # Convert results to a tuple of id+info and a record of the artifacts
342 # that should not be deleted from datastore. The id+info tuple is
343 # solely to allow logging to report the relevant ID.
344 id_info = ((FakeDatasetRef(row["dataset_id"]), record_class.from_record(row)) for row in rows)
346 # Start contextmanager, return results
347 yield ((id_info, preserved))
349 # No exception raised in context manager block.
350 if not rows or dry_run:
351 return
353 # Delete the rows from the records table
354 records_table.delete(["dataset_id"], *[{"dataset_id": row["dataset_id"]} for row in rows])
356 # Delete those rows from the trash table.
357 self._db.delete(
358 self._tables.dataset_location_trash,
359 ["dataset_id", "datastore_name"],
360 *[{"dataset_id": row["dataset_id"], "datastore_name": row["datastore_name"]} for row in rows],
361 )
364class MonolithicDatastoreRegistryBridgeManager(DatastoreRegistryBridgeManager):
365 """An implementation of `DatastoreRegistryBridgeManager` that uses the same
366 two tables for all non-ephemeral datastores.
368 Parameters
369 ----------
370 db : `Database`
371 Object providing a database connection and generic distractions.
372 tables : `_TablesTuple`
373 Named tuple containing `sqlalchemy.schema.Table` instances.
374 opaque : `OpaqueTableStorageManager`
375 Manager object for opaque table storage in the `Registry`.
376 universe : `DimensionUniverse`
377 All dimensions know to the `Registry`.
378 registry_schema_version : `VersionTuple` or `None`, optional
379 The version of the registry schema.
380 """
382 def __init__(
383 self,
384 *,
385 db: Database,
386 tables: _TablesTuple,
387 opaque: OpaqueTableStorageManager,
388 universe: DimensionUniverse,
389 registry_schema_version: VersionTuple | None = None,
390 ):
391 super().__init__(
392 opaque=opaque,
393 universe=universe,
394 registry_schema_version=registry_schema_version,
395 )
396 self._db = db
397 self._tables = tables
398 self._ephemeral: dict[str, EphemeralDatastoreRegistryBridge] = {}
400 def clone(self, *, db: Database, opaque: OpaqueTableStorageManager) -> DatastoreRegistryBridgeManager:
401 return MonolithicDatastoreRegistryBridgeManager(
402 db=db,
403 tables=self._tables,
404 opaque=opaque,
405 universe=self.universe,
406 registry_schema_version=self._registry_schema_version,
407 )
409 @classmethod
410 def initialize(
411 cls,
412 db: Database,
413 context: StaticTablesContext,
414 *,
415 opaque: OpaqueTableStorageManager,
416 datasets: type[DatasetRecordStorageManager],
417 universe: DimensionUniverse,
418 registry_schema_version: VersionTuple | None = None,
419 ) -> DatastoreRegistryBridgeManager:
420 # Docstring inherited from DatastoreRegistryBridge
421 tables = context.addTableTuple(_makeTableSpecs(datasets))
422 return cls(
423 db=db,
424 tables=cast(_TablesTuple, tables),
425 opaque=opaque,
426 universe=universe,
427 registry_schema_version=registry_schema_version,
428 )
430 def refresh(self) -> None:
431 # Docstring inherited from DatastoreRegistryBridge
432 # This implementation has no in-Python state that depends on which
433 # datastores exist, so there's nothing to do.
434 pass
436 def register(self, name: str, *, ephemeral: bool = False) -> DatastoreRegistryBridge:
437 # Docstring inherited from DatastoreRegistryBridge
438 if ephemeral: 438 ↛ 439line 438 didn't jump to line 439 because the condition on line 438 was never true
439 return self._ephemeral.setdefault(name, EphemeralDatastoreRegistryBridge(name))
440 return MonolithicDatastoreRegistryBridge(name, db=self._db, tables=self._tables)
442 def findDatastores(self, ref: DatasetIdRef) -> Iterable[str]:
443 # Docstring inherited from DatastoreRegistryBridge
444 sql = (
445 sqlalchemy.sql.select(self._tables.dataset_location.columns.datastore_name)
446 .select_from(self._tables.dataset_location)
447 .where(self._tables.dataset_location.columns.dataset_id == ref.id)
448 )
449 with self._db.query(sql) as sql_result:
450 sql_rows = sql_result.mappings().fetchall()
451 for row in sql_rows:
452 yield row[self._tables.dataset_location.columns.datastore_name]
453 for name, bridge in self._ephemeral.items():
454 if ref in bridge:
455 yield name
457 @classmethod
458 def currentVersions(cls) -> list[VersionTuple]:
459 # Docstring inherited from VersionedExtension.
460 return [_VERSION]