Remove daf_relation from registry.certify()

Modify registry.certify() to use the new query system instead of daf_relation, because we plan to retire daf_relation.

Author: dhirving

python/lsst/daf/butler/registry/datasets/byDimensions/_manager.py

@@ -5,7 +5,8 @@
 import dataclasses
 import datetime
 import logging
-from collections.abc import Iterable, Mapping, Sequence, Set
+from collections.abc import Callable, Iterable, Mapping, Sequence, Set
+from contextlib import AbstractContextManager
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import astropy.time
@@ -24,6 +25,7 @@
 from ...._timespan import Timespan
 from ....dimensions import DataCoordinate, DimensionGroup, DimensionUniverse
 from ....direct_query_driver import SqlJoinsBuilder, SqlSelectBuilder  # new query system, server+direct only
+from ....queries import Query
 from ....queries import tree as qt  # new query system, both clients + server
 from ..._caching_context import CachingContext
 from ..._collection_summary import CollectionSummary
@@ -966,7 +968,7 @@ def certify(
         collection: CollectionRecord,
         datasets: Iterable[DatasetRef],
         timespan: Timespan,
-        context: SqlQueryContext,
+        query_func: Callable[[], AbstractContextManager[Query]],
     ) -> None:
         # Docstring inherited from DatasetRecordStorageManager.
         if (storage := self._find_storage(dataset_type.name)) is None:
@@ -1019,23 +1021,25 @@ def certify(
                 ) from err
         else:
             # Have to implement exclusion constraint ourselves.
-            # Start by building a SELECT query for any rows that would overlap
-            # this one.
-            relation = self._build_calib_overlap_query(dataset_type, collection, data_ids, timespan, context)
             # Acquire a table lock to ensure there are no concurrent writes
             # could invalidate our checking before we finish the inserts.  We
             # use a SAVEPOINT in case there is an outer transaction that a
             # failure here should not roll back.
-            with self._db.transaction(lock=[calibs_table], savepoint=True):
-                # Enter SqlQueryContext in case we need to use a temporary
-                # table to include the give data IDs in the query.  Note that
-                # by doing this inside the transaction, we make sure it doesn't
-                # attempt to close the session when its done, since it just
-                # sees an already-open session that it knows it shouldn't
-                # manage.
-                with context:
-                    # Run the check SELECT query.
-                    conflicting = context.count(context.process(relation))
+            with self._db.transaction(
+                lock=[calibs_table],
+                savepoint=True,
+                # join_data_coordinates sometimes requires a temp table
+                for_temp_tables=True,
+            ):
+                # Query for any rows that would overlap this one.
+                with query_func() as query:
+                    if data_ids is not None:
+                        query = query.join_data_coordinates(data_ids)
+                    timespan_column = query.expression_factory[dataset_type.name].timespan
+                    result = query.datasets(dataset_type, collection.name).where(
+                        timespan_column.overlaps(timespan)
+                    )
+                    conflicting = result.count()
                     if conflicting > 0:
                         raise ConflictingDefinitionError(
                             f"{conflicting} validity range conflicts certifying datasets of type "

python/lsst/daf/butler/registry/interfaces/_datasets.py

@@ -32,7 +32,8 @@
 __all__ = ("DatasetRecordStorageManager",)
 
 from abc import abstractmethod
-from collections.abc import Iterable, Mapping, Sequence, Set
+from collections.abc import Callable, Iterable, Mapping, Sequence, Set
+from contextlib import AbstractContextManager
 from typing import TYPE_CHECKING, Any
 
 from lsst.daf.relation import Relation
@@ -42,6 +43,7 @@
 from ..._exceptions import DatasetTypeError, DatasetTypeNotSupportedError
 from ..._timespan import Timespan
 from ...dimensions import DataCoordinate
+from ...queries import Query
 from ._versioning import VersionedExtension, VersionTuple
 
 if TYPE_CHECKING:
@@ -504,7 +506,7 @@ def certify(
         collection: CollectionRecord,
         datasets: Iterable[DatasetRef],
         timespan: Timespan,
-        context: SqlQueryContext,
+        query_func: Callable[[], AbstractContextManager[Query]],
     ) -> None:
         """Associate one or more datasets with a calibration collection and a
         validity range within it.
@@ -521,9 +523,10 @@ def certify(
             `DatasetType` as ``dataset_type``, but this is not checked.
         timespan : `Timespan`
             The validity range for these datasets within the collection.
-        context : `SqlQueryContext`
-            The object that manages database connections, temporary tables and
-            relation engines for this query.
+        query_func : `Callable` [[], `AbstractContextManager` [ `Query` ]],
+            Function returning a context manager that sets up a `Query` object
+            for querying the registry. (That is, a function equivalent to
+            `Butler.query()`).
 
         Raises
         ------

python/lsst/daf/butler/registry/sql_registry.py

@@ -1161,17 +1161,14 @@ def certify(self, collection: str, refs: Iterable[DatasetRef], timespan: Timespa
             ``collection.type is not CollectionType.CALIBRATION``.
         """
         progress = Progress("lsst.daf.butler.Registry.certify", level=logging.DEBUG)
-        collectionRecord = self._managers.collections.find(collection)
-        for datasetType, refsForType in progress.iter_item_chunks(
-            DatasetRef.iter_by_type(refs), desc="Certifying datasets by type"
-        ):
-            self._managers.datasets.certify(
-                datasetType,
-                collectionRecord,
-                refsForType,
-                timespan,
-                context=queries.SqlQueryContext(self._db, self._managers.column_types),
-            )
+        with self._managers.caching_context.enable_collection_record_cache():
+            collectionRecord = self._managers.collections.find(collection)
+            for datasetType, refsForType in progress.iter_item_chunks(
+                DatasetRef.iter_by_type(refs), desc="Certifying datasets by type"
+            ):
+                self._managers.datasets.certify(
+                    datasetType, collectionRecord, refsForType, timespan, self._query
+                )
 
     @transactional
     def decertify(