DM-37173: Rewrite getURIs to use getManyURIs

python/lsst/daf/butler/datastores/fileDatastore.py

@@ -369,7 +369,7 @@ def _delete_artifact(self, location: Location) -> None:
     def addStoredItemInfo(self, refs: Iterable[DatasetRef], infos: Iterable[StoredFileInfo]) -> None:
         # Docstring inherited from GenericBaseDatastore
         records = [info.rebase(ref).to_record() for ref, info in zip(refs, infos)]
-        self._table.insert(*records)
+        self._table.insert(*records, transaction=self._transaction)
 
     def getStoredItemsInfo(self, ref: DatasetIdRef) -> List[StoredFileInfo]:
         # Docstring inherited from GenericBaseDatastore
@@ -2838,7 +2838,7 @@ def import_records(self, data: Mapping[str, DatastoreRecordData]) -> None:
                     assert isinstance(info, StoredFileInfo), "Expecting StoredFileInfo records"
                     unpacked_records.append(info.to_record())
         if unpacked_records:
-            self._table.insert(*unpacked_records)
+            self._table.insert(*unpacked_records, transaction=self._transaction)
 
     def export_records(self, refs: Iterable[DatasetIdRef]) -> Mapping[str, DatastoreRecordData]:
         # Docstring inherited from the base class.

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

@@ -27,12 +27,15 @@
 __all__ = ["OpaqueTableStorageManager", "OpaqueTableStorage"]
 
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, Iterator, Optional
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Optional
 
 from ...core.ddl import TableSpec
 from ._database import Database, StaticTablesContext
 from ._versioning import VersionedExtension
 
+if TYPE_CHECKING:
+    from ...core.datastore import DatastoreTransaction
+
 
 class OpaqueTableStorage(ABC):
     """An interface that manages the records associated with a particular
@@ -48,14 +51,17 @@ def __init__(self, name: str):
         self.name = name
 
     @abstractmethod
-    def insert(self, *data: dict) -> None:
+    def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
         """Insert records into the table
 
         Parameters
         ----------
         *data
             Each additional positional argument is a dictionary that represents
             a single row to be added.
+        transaction : `DatastoreTransaction`, optional
+            Transaction object. Can be `None` in some bridges if no rollback
+            is required.
         """
         raise NotImplementedError()
 

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

@@ -27,7 +27,7 @@
 __all__ = ["ByNameOpaqueTableStorage", "ByNameOpaqueTableStorageManager"]
 
 import itertools
-from typing import Any, ClassVar, Dict, Iterable, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Iterable, Iterator, List, Optional
 
 import sqlalchemy
 
@@ -40,6 +40,9 @@
     VersionTuple,
 )
 
+if TYPE_CHECKING:
+    from ..core.datastore import DatastoreTransaction
+
 # This has to be updated on every schema change
 _VERSION = VersionTuple(0, 2, 0)
 
@@ -68,7 +71,7 @@ def __init__(self, *, db: Database, name: str, table: sqlalchemy.schema.Table):
         self._db = db
         self._table = table
 
-    def insert(self, *data: dict) -> None:
+    def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
         # Docstring inherited from OpaqueTableStorage.
         self._db.insert(self._table, *data)
 

python/lsst/daf/butler/tests/_dummyRegistry.py

@@ -40,14 +40,16 @@
     VersionTuple,
 )
 
+from ..core.datastore import DatastoreTransaction
+
 
 class DummyOpaqueTableStorage(OpaqueTableStorage):
     def __init__(self, name: str, spec: ddl.TableSpec) -> None:
         super().__init__(name=name)
         self._rows: list[dict] = []
         self._spec = spec
 
-    def insert(self, *data: dict) -> None:
+    def insert(self, *data: dict, transaction: DatastoreTransaction | None = None) -> None:
         # Docstring inherited from OpaqueTableStorage.
         uniqueConstraints = list(self._spec.unique)
         uniqueConstraints.append(tuple(field.name for field in self._spec.fields if field.primaryKey))
@@ -59,6 +61,8 @@ def insert(self, *data: dict) -> None:
                         f"Unique constraint {constraint} violation in external table {self.name}."
                     )
             self._rows.append(d)
+            if transaction is not None:
+                transaction.registerUndo("insert", self.delete, [], d)
 
     def fetch(self, **where: Any) -> Iterator[dict]:
         # Docstring inherited from OpaqueTableStorage.