Route internal metastore queries through the select seams (#1889)

Author: amritghimire

src/datachain/data_storage/metastore.py

@@ -1080,7 +1080,9 @@ def remove_namespace(self, namespace_id: int) -> None:
             # sneak in between the count and the delete below. No-op on
             # SQLite (writes are already serialized at the database level).
             self.db.execute(
-                select(n.c.id).where(n.c.id == namespace_id).with_for_update()
+                self._namespaces_select(n.c.id)
+                .where(n.c.id == namespace_id)
+                .with_for_update()
             )
             num_projects = self.count_projects(namespace_id)
             if num_projects > 0:
@@ -1216,7 +1218,9 @@ def remove_project(self, project_id: int) -> None:
             # sneak in between the count and the delete below. No-op on
             # SQLite (writes are already serialized at the database level).
             self.db.execute(
-                select(p.c.id).where(p.c.id == project_id).with_for_update()
+                self._projects_select(p.c.id)
+                .where(p.c.id == project_id)
+                .with_for_update()
             )
             num_datasets = self.count_datasets(project_id)
             if num_datasets > 0:
@@ -1862,7 +1866,9 @@ def remove_dataset_version(
             # if-empty below. No-op on SQLite (writes are already
             # serialized at the database level).
             self.db.execute(
-                select(d.c.id).where(d.c.id == dataset.id).with_for_update()
+                self._datasets_select(d.c.id)
+                .where(d.c.id == dataset.id)
+                .with_for_update()
             )
 
             self.db.execute(
@@ -1874,7 +1880,7 @@ def remove_dataset_version(
             # Count from DB - in-memory dataset.versions may be incomplete.
             remaining = next(
                 self.db.execute(
-                    select(f.count())
+                    self._datasets_versions_select(f.count())
                     .select_from(dv)
                     .where(dv.c.dataset_id == dataset.id)
                 )
@@ -2188,6 +2194,16 @@ def update_dataset_dependency_source(
         q = q.values(**data)
         self.db.execute(q)
 
+    @cached_property
+    def _dependency_datasets_source(self):
+        # Source alias for direct-dependency queries; cached so the column hook
+        # and the FROM clause reference the same object. Override to filter.
+        return self._datasets
+
+    @cached_property
+    def _dependency_namespaces_source(self):
+        return self._namespaces
+
     @abstractmethod
     def _dataset_dependencies_select_columns(self) -> list["SchemaItem"]:
         """
@@ -2211,9 +2227,9 @@ def get_direct_dataset_dependencies(
         dataset: DatasetRecord,
         version: str,
     ) -> list[DatasetDependency | None]:
-        n = self._namespaces
+        n = self._dependency_namespaces_source
         p = self._projects
-        d = self._datasets
+        d = self._dependency_datasets_source
         dd = self._datasets_dependencies
         dv = self._datasets_versions
 
@@ -2925,7 +2941,7 @@ def link_dataset_version_to_job(
 
             # Also update dataset_version.job_id to point to this job
             update_query = (
-                self._datasets_versions.update()
+                self._datasets_versions_update()
                 .where(self._datasets_versions.c.id == dataset_version_id)
                 .values(job_id=job_id)
             )