Refactor dataset dependency query to make select column configurable (#1405)

* Refactor dataset dependency query to make select column configurable

# Changes:
- Added new abstract method `_dataset_dependency_nodes_select_columns()` to make the column selection for dataset dependency queries more maintainable and extensible across different database backends
- Introduced a depth limit of 100 in the recursive CTE to prevent infinite loops

* Move to constant

Author: amritghimire

src/datachain/data_storage/metastore.py

@@ -56,13 +56,15 @@
 from datachain.utils import JSONSerialize
 
 if TYPE_CHECKING:
-    from sqlalchemy import Delete, Insert, Select, Update
+    from sqlalchemy import CTE, Delete, Insert, Select, Subquery, Update
     from sqlalchemy.schema import SchemaItem
+    from sqlalchemy.sql.elements import ColumnElement
 
     from datachain.data_storage import schema
     from datachain.data_storage.db_engine import DatabaseEngine
 
 logger = logging.getLogger("datachain")
+DEPTH_LIMIT_DEFAULT = 100
 
 
 class AbstractMetastore(ABC, Serializable):
@@ -1463,6 +1465,18 @@ def _dataset_dependencies_select_columns(self) -> list["SchemaItem"]:
         Returns a list of columns to select in a query for fetching dataset dependencies
         """
 
+    @abstractmethod
+    def _dataset_dependency_nodes_select_columns(
+        self,
+        namespaces_subquery: "Subquery",
+        dependency_tree_cte: "CTE",
+        datasets_subquery: "Subquery",
+    ) -> list["ColumnElement"]:
+        """
+        Returns a list of columns to select in a query for fetching
+        dataset dependency nodes.
+        """
+
     def get_direct_dataset_dependencies(
         self, dataset: DatasetRecord, version: str
     ) -> list[DatasetDependency | None]:
@@ -1493,7 +1507,7 @@ def get_direct_dataset_dependencies(
         return [self.dependency_class.parse(*r) for r in self.db.execute(query)]
 
     def get_dataset_dependency_nodes(
-        self, dataset_id: int, version_id: int
+        self, dataset_id: int, version_id: int, depth_limit: int = DEPTH_LIMIT_DEFAULT
     ) -> list[DatasetDependencyNode | None]:
         n = self._namespaces_select().subquery()
         p = self._projects
@@ -1522,33 +1536,31 @@ def get_dataset_dependency_nodes(
         cte = base_query.cte(name="dependency_tree", recursive=True)
 
         # Recursive case: dependencies of dependencies
-        recursive_query = select(
-            *dep_fields,
-            (cte.c.depth + 1).label("depth"),
-        ).select_from(
-            cte.join(
-                dd,
-                (cte.c.dataset_id == dd.c.source_dataset_id)
-                & (cte.c.dataset_version_id == dd.c.source_dataset_version_id),
+        # Limit depth to 100 to prevent infinite loops in case of circular dependencies
+        recursive_query = (
+            select(
+                *dep_fields,
+                (cte.c.depth + 1).label("depth"),
             )
+            .select_from(
+                cte.join(
+                    dd,
+                    (cte.c.dataset_id == dd.c.source_dataset_id)
+                    & (cte.c.dataset_version_id == dd.c.source_dataset_version_id),
+                )
+            )
+            .where(cte.c.depth < depth_limit)
         )
 
         cte = cte.union(recursive_query)
 
         # Fetch all with full details
-        final_query = select(
-            n.c.name,
-            p.c.name,
-            cte.c.id,
-            cte.c.dataset_id,
-            cte.c.dataset_version_id,
-            d.c.name,
-            dv.c.version,
-            dv.c.created_at,
-            cte.c.source_dataset_id,
-            cte.c.source_dataset_version_id,
-            cte.c.depth,
-        ).select_from(
+        select_cols = self._dataset_dependency_nodes_select_columns(
+            namespaces_subquery=n,
+            dependency_tree_cte=cte,
+            datasets_subquery=d,
+        )
+        final_query = self._datasets_dependencies_select(*select_cols).select_from(
             # Use outer joins to handle cases where dependent datasets have been
             # physically deleted. This allows us to return dependency records with
             # None values instead of silently omitting them, making broken

src/datachain/data_storage/sqlite.py

@@ -20,7 +20,10 @@
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
 from sqlalchemy.sql import func
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from sqlalchemy.sql.elements import (
+    BinaryExpression,
+    BooleanClauseList,
+)
 from sqlalchemy.sql.expression import bindparam, cast
 from sqlalchemy.sql.selectable import Select
 from tqdm.auto import tqdm
@@ -41,6 +44,7 @@
 from datachain.utils import DataChainDir, batched, batched_it
 
 if TYPE_CHECKING:
+    from sqlalchemy import CTE, Subquery
     from sqlalchemy.dialects.sqlite import Insert
     from sqlalchemy.engine.base import Engine
     from sqlalchemy.schema import SchemaItem
@@ -539,6 +543,26 @@ def _dataset_dependencies_select_columns(self) -> list["SchemaItem"]:
             self._datasets_versions.c.created_at,
         ]
 
+    def _dataset_dependency_nodes_select_columns(
+        self,
+        namespaces_subquery: "Subquery",
+        dependency_tree_cte: "CTE",
+        datasets_subquery: "Subquery",
+    ) -> list["ColumnElement"]:
+        return [
+            namespaces_subquery.c.name,
+            self._projects.c.name,
+            dependency_tree_cte.c.id,
+            dependency_tree_cte.c.dataset_id,
+            dependency_tree_cte.c.dataset_version_id,
+            datasets_subquery.c.name,
+            self._datasets_versions.c.version,
+            self._datasets_versions.c.created_at,
+            dependency_tree_cte.c.source_dataset_id,
+            dependency_tree_cte.c.source_dataset_version_id,
+            dependency_tree_cte.c.depth,
+        ]
+
     #
     # Jobs
     #