DM-36172: Add explicit transaction around moveToTrash calls

doc/changes/DM-36172.misc.rst

@@ -0,0 +1,2 @@
+Add support for in-memory datastore to roll back a call to ``datastore.trash()``.
+This required that the ``bridge.moveToTrash()`` method now takes an additional ``transaction`` parameter (that can be `None`).

python/lsst/daf/butler/_butler.py

@@ -1636,13 +1636,14 @@ def removeRuns(self, names: Iterable[str], unstore: bool = True) -> None:
             if collectionType is not CollectionType.RUN:
                 raise TypeError(f"The collection type of '{name}' is {collectionType.name}, not RUN.")
             refs.extend(self.registry.queryDatasets(..., collections=name, findFirst=True))
-        with self.registry.transaction():
-            if unstore:
-                self.datastore.trash(refs)
-            else:
-                self.datastore.forget(refs)
-            for name in names:
-                self.registry.removeCollection(name)
+        with self.datastore.transaction():
+            with self.registry.transaction():
+                if unstore:
+                    self.datastore.trash(refs)
+                else:
+                    self.datastore.forget(refs)
+                for name in names:
+                    self.registry.removeCollection(name)
         if unstore:
             # Point of no return for removing artifacts
             self.datastore.emptyTrash()
@@ -1703,14 +1704,15 @@ def remove(child: str, parent: str) -> None:
                 raise RuntimeError(f"{name} is not a child of {parent}") from e
             self.registry.setCollectionChain(parent, chain)
 
-        with self.registry.transaction():
-            if unlink:
-                for parent in unlink:
-                    remove(name, parent)
-            if unstore:
-                refs = self.registry.queryDatasets(..., collections=name, findFirst=True)
-                self.datastore.trash(refs)
-            self.registry.removeCollection(name)
+        with self.datastore.transaction():
+            with self.registry.transaction():
+                if unlink:
+                    for parent in unlink:
+                        remove(name, parent)
+                if unstore:
+                    refs = self.registry.queryDatasets(..., collections=name, findFirst=True)
+                    self.datastore.trash(refs)
+                self.registry.removeCollection(name)
 
         if unstore:
             # Point of no return for removing artifacts
@@ -1765,15 +1767,16 @@ def pruneDatasets(
         # mutating the Registry (it can _look_ at Datastore-specific things,
         # but shouldn't change them), and hence all operations here are
         # Registry operations.
-        with self.registry.transaction():
-            if unstore:
-                self.datastore.trash(refs)
-            if purge:
-                self.registry.removeDatasets(refs)
-            elif disassociate:
-                assert tags, "Guaranteed by earlier logic in this function."
-                for tag in tags:
-                    self.registry.disassociate(tag, refs)
+        with self.datastore.transaction():
+            with self.registry.transaction():
+                if unstore:
+                    self.datastore.trash(refs)
+                if purge:
+                    self.registry.removeDatasets(refs)
+                elif disassociate:
+                    assert tags, "Guaranteed by earlier logic in this function."
+                    for tag in tags:
+                        self.registry.disassociate(tag, refs)
         # We've exited the Registry transaction, and apparently committed.
         # (if there was an exception, everything rolled back, and it's as if
         # nothing happened - and we never get here).

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

@@ -2246,7 +2246,7 @@ def trash(self, ref: Union[DatasetRef, Iterable[DatasetRef]], ignore_errors: boo
             log.debug("Doing multi-dataset trash in datastore %s", self.name)
             # Assumed to be an iterable of refs so bulk mode enabled.
             try:
-                self.bridge.moveToTrash(ref)
+                self.bridge.moveToTrash(ref, transaction=self._transaction)
             except Exception as e:
                 if ignore_errors:
                     log.warning("Unexpected issue moving multiple datasets to trash: %s", e)
@@ -2280,7 +2280,7 @@ def trash(self, ref: Union[DatasetRef, Iterable[DatasetRef]], ignore_errors: boo
 
         # Mark dataset as trashed
         try:
-            self.bridge.moveToTrash([ref])
+            self.bridge.moveToTrash([ref], transaction=self._transaction)
         except Exception as e:
             if ignore_errors:
                 log.warning(

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

@@ -39,6 +39,7 @@
     StorageClass,
     StoredDatastoreItemInfo,
 )
+from lsst.daf.butler.core.utils import transactional
 from lsst.daf.butler.registry.interfaces import DatastoreRegistryBridge
 from lsst.resources import ResourcePath
 
@@ -512,6 +513,7 @@ def forget(self, refs: Iterable[DatasetRef]) -> None:
         for ref in refs:
             self.removeStoredItemInfo(ref)
 
+    @transactional
     def trash(self, ref: Union[DatasetRef, Iterable[DatasetRef]], ignore_errors: bool = False) -> None:
         """Indicate to the Datastore that a dataset can be removed.
 
@@ -536,7 +538,7 @@ def trash(self, ref: Union[DatasetRef, Iterable[DatasetRef]], ignore_errors: boo
         """
         if not isinstance(ref, DatasetRef):
             log.debug("Bulk trashing of datasets in datastore %s", self.name)
-            self.bridge.moveToTrash(ref)
+            self.bridge.moveToTrash(ref, transaction=self._transaction)
             return
 
         log.debug("Trash %s in datastore %s", ref, self.name)
@@ -546,7 +548,7 @@ def trash(self, ref: Union[DatasetRef, Iterable[DatasetRef]], ignore_errors: boo
             self._get_dataset_info(ref)
 
             # Move datasets to trash table
-            self.bridge.moveToTrash([ref])
+            self.bridge.moveToTrash([ref], transaction=self._transaction)
         except Exception as e:
             if ignore_errors:
                 log.warning(

python/lsst/daf/butler/registry/bridge/ephemeral.py

@@ -30,6 +30,7 @@
 
 if TYPE_CHECKING:
     from ...core import StoredDatastoreItemInfo
+    from ...core.datastore import DatastoreTransaction
 
 
 class EphemeralDatastoreRegistryBridge(DatastoreRegistryBridge):
@@ -62,9 +63,18 @@ def insert(self, refs: Iterable[DatasetIdRef]) -> None:
     def forget(self, refs: Iterable[DatasetIdRef]) -> None:
         self._datasetIds.difference_update(ref.id for ref in refs)
 
-    def moveToTrash(self, refs: Iterable[DatasetIdRef]) -> None:
+    def _rollbackMoveToTrash(self, refs: Iterable[DatasetIdRef]) -> None:
+        """Rollback a moveToTrash call."""
+        for ref in refs:
+            self._trashedIds.remove(ref.getCheckedId())
+
+    def moveToTrash(self, refs: Iterable[DatasetIdRef], transaction: Optional[DatastoreTransaction]) -> None:
         # Docstring inherited from DatastoreRegistryBridge
-        self._trashedIds.update(ref.getCheckedId() for ref in refs)
+        if transaction is None:
+            raise RuntimeError("Must be called with a defined transaction.")
+        ref_list = list(refs)
+        with transaction.undoWith(f"Trash {len(ref_list)} datasets", self._rollbackMoveToTrash, ref_list):
+            self._trashedIds.update(ref.getCheckedId() for ref in ref_list)
 
     def check(self, refs: Iterable[DatasetIdRef]) -> Iterable[DatasetIdRef]:
         # Docstring inherited from DatastoreRegistryBridge

python/lsst/daf/butler/registry/bridge/monolithic.py

@@ -43,6 +43,7 @@
 
 if TYPE_CHECKING:
     from ...core import DimensionUniverse
+    from ...core.datastore import DatastoreTransaction
     from ..interfaces import (
         Database,
         DatasetRecordStorageManager,
@@ -157,7 +158,7 @@ def forget(self, refs: Iterable[DatasetIdRef]) -> None:
         rows = self._refsToRows(self.check(refs))
         self._db.delete(self._tables.dataset_location, ["datastore_name", "dataset_id"], *rows)
 
-    def moveToTrash(self, refs: Iterable[DatasetIdRef]) -> None:
+    def moveToTrash(self, refs: Iterable[DatasetIdRef], transaction: Optional[DatastoreTransaction]) -> None:
         # Docstring inherited from DatastoreRegistryBridge
         # TODO: avoid self.check() call via queries like
         #     INSERT INTO dataset_location_trash

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

@@ -32,6 +32,7 @@
 
 if TYPE_CHECKING:
     from ...core import DatasetType, DimensionUniverse, StoredDatastoreItemInfo
+    from ...core.datastore import DatastoreTransaction
     from ._database import Database, StaticTablesContext
     from ._datasets import DatasetRecordStorageManager
     from ._opaque import OpaqueTableStorage, OpaqueTableStorageManager
@@ -150,13 +151,16 @@ def forget(self, refs: Iterable[DatasetIdRef]) -> None:
         raise NotImplementedError()
 
     @abstractmethod
-    def moveToTrash(self, refs: Iterable[DatasetIdRef]) -> None:
+    def moveToTrash(self, refs: Iterable[DatasetIdRef], transaction: Optional[DatastoreTransaction]) -> None:
         """Move dataset location information to trash.
 
         Parameters
         ----------
         refs : `Iterable` of `DatasetIdRef`
             References to the datasets.
+        transaction : `DatastoreTransaction` or `None`
+            Transaction object. Can be `None` in some bridges or if no rollback
+            is required.
 
         Raises
         ------

tests/test_butler.py

@@ -799,10 +799,10 @@ def testPruneCollections(self):
         # Try to delete RUN collections, which should fail with complete
         # rollback because they're still referenced by the CHAINED
         # collection.
-        with self.assertRaises(Exception):
-            butler.pruneCollection(run1, pruge=True, unstore=True)
-        with self.assertRaises(Exception):
-            butler.pruneCollection(run2, pruge=True, unstore=True)
+        with self.assertRaises(sqlalchemy.exc.IntegrityError):
+            butler.pruneCollection(run1, purge=True, unstore=True)
+        with self.assertRaises(sqlalchemy.exc.IntegrityError):
+            butler.pruneCollection(run2, purge=True, unstore=True)
         self.assertCountEqual(set(butler.registry.queryDatasets(..., collections=...)), [ref1, ref2, ref3])
         existence = butler.datastore.mexists([ref1, ref2, ref3])
         self.assertTrue(existence[ref1])
@@ -1388,7 +1388,7 @@ def testPruneDatasets(self):
         # Check that in normal mode, deleting the record will lead to
         # trash not touching the file.
         uri1 = butler.datastore.getURI(ref1)
-        butler.datastore.bridge.moveToTrash([ref1])  # Update the dataset_location table
+        butler.datastore.bridge.moveToTrash([ref1], transaction=None)  # Update the dataset_location table
         butler.datastore._table.delete(["dataset_id"], {"dataset_id": ref1.id})
         butler.datastore.trash(ref1)
         butler.datastore.emptyTrash()
@@ -1404,7 +1404,7 @@ def testPruneDatasets(self):
         self.assertTrue(uri3.exists())
 
         # Remove the datastore record.
-        butler.datastore.bridge.moveToTrash([ref2])  # Update the dataset_location table
+        butler.datastore.bridge.moveToTrash([ref2], transaction=None)  # Update the dataset_location table
         butler.datastore._table.delete(["dataset_id"], {"dataset_id": ref2.id})
         self.assertTrue(uri2.exists())
         butler.datastore.trash([ref2, ref3])

ups/daf_butler.table

@@ -5,4 +5,3 @@ setupRequired(resources)
 
 envPrepend(PATH, ${PRODUCT_DIR}/bin)
 envPrepend(PYTHONPATH, ${PRODUCT_DIR}/python)
-envPrepend(MYPYPATH, ${PRODUCT_DIR}/python)