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