track remote tiledbsoma stores for writes

lamindb/_artifact.py

@@ -879,16 +879,32 @@ def open(
             f" {', '.join(suffixes[:-1])}."
         )
 
-    from lamindb.core.storage._backed_access import backed_access
+    from lamindb.core.storage._backed_access import _track_writes_factory, backed_access
 
     using_key = settings._using_key
     filepath = filepath_from_artifact(self, using_key=using_key)
+    is_tiledbsoma = filepath.name == "soma" or filepath.suffix == ".tiledbsoma"
     # consider the case where an object is already locally cached
     localpath = setup_settings.instance.storage.cloud_to_local_no_update(filepath)
-    if localpath.exists():
+    if not is_tiledbsoma and localpath.exists():
         access = backed_access(localpath, mode, using_key)
     else:
         access = backed_access(filepath, mode, using_key)
+        if is_tiledbsoma and mode == "w":
+            if access.__class__.__name__ not in {"Collection", "Experiment"}:
+                raise ValueError(
+                    "The store seems to be tiledbsoma but it is neither `Collection` nor `Experiment`."
+                )
+
+            def finalize():
+                nonlocal self, filepath
+                _, hash, _, _ = get_stat_dir_cloud(filepath)
+                if self.hash != hash:
+                    logger.warning(
+                        "The hash of the tiledbsoma store has changed, consider updating this artifact."
+                    )
+
+            access = _track_writes_factory(access, finalize)
     # only call if open is successfull
     _track_run_input(self, is_run_input)
     return access

lamindb/core/storage/_backed_access.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Callable, Literal
 
 from anndata._io.specs.registry import get_spec
 from lnschema_core import Artifact
@@ -16,6 +16,39 @@
     from upath import UPath
 
 
+def _track_writes_factory(obj: Any, finalize: Callable):
+    closed: bool = False
+
+    tracked_class = obj.__class__
+    type_dict = {"__doc__": tracked_class.__doc__}
+    if hasattr(tracked_class, "__slots__"):
+        type_dict["__slots__"] = ()
+    if hasattr(tracked_class, "__exit__"):
+
+        def __exit__(self, exc_type, exc_val, exc_tb):
+            nonlocal closed
+            tracked_class.__exit__(self, exc_type, exc_val, exc_tb)
+            if not closed:
+                finalize()
+                closed = True
+
+        type_dict["__exit__"] = __exit__
+    if hasattr(tracked_class, "close"):
+
+        def close(self, *args, **kwargs):
+            nonlocal closed
+            tracked_class.close(self, *args, **kwargs)
+            if not closed:
+                finalize()
+                closed = True
+
+        type_dict["close"] = close
+
+    Track = type(tracked_class.__name__ + "Track", (tracked_class,), type_dict)
+    obj.__class__ = Track
+    return obj
+
+
 def _open_tiledbsoma(
     filepath: UPath, mode: Literal["r", "w"] = "r"
 ) -> SOMACollection | SOMAExperiment: