datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/data_storage/metastore.py

@@ -1,28 +1,37 @@
 import copy
-import json
 import logging
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from contextlib import contextmanager, suppress
 from datetime import datetime, timezone
 from functools import cached_property, reduce
 from itertools import groupby
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 from sqlalchemy import (
     JSON,
     BigInteger,
+    Boolean,
     Column,
     DateTime,
     ForeignKey,
+    Index,
     Integer,
     Table,
     Text,
     UniqueConstraint,
+    cast,
+    desc,
+    literal,
     select,
 )
+from sqlalchemy.sql import func as f
 
+from datachain import json
+from datachain.catalog.dependency import DatasetDependencyNode
+from datachain.checkpoint import Checkpoint
 from datachain.data_storage import JobQueryType, JobStatus
 from datachain.data_storage.serializer import Serializable
 from datachain.dataset import (
@@ -33,27 +42,34 @@ from datachain.dataset import (
     DatasetStatus,
     DatasetVersion,
     StorageURI,
+    parse_schema,
 )
 from datachain.error import (
+    CheckpointNotFoundError,
+    DataChainError,
     DatasetNotFoundError,
     DatasetVersionNotFoundError,
+    NamespaceDeleteNotAllowedError,
     NamespaceNotFoundError,
+    ProjectDeleteNotAllowedError,
     ProjectNotFoundError,
     TableMissingError,
 )
 from datachain.job import Job
 from datachain.namespace import Namespace
 from datachain.project import Project
-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
+JOB_ANCESTRY_MAX_DEPTH = 100
 
 
 class AbstractMetastore(ABC, Serializable):
@@ -68,14 +84,17 @@ class AbstractMetastore(ABC, Serializable):
     namespace_class: type[Namespace] = Namespace
     project_class: type[Project] = Project
     dataset_class: type[DatasetRecord] = DatasetRecord
+    dataset_version_class: type[DatasetVersion] = DatasetVersion
     dataset_list_class: type[DatasetListRecord] = DatasetListRecord
     dataset_list_version_class: type[DatasetListVersion] = DatasetListVersion
     dependency_class: type[DatasetDependency] = DatasetDependency
+    dependency_node_class: type[DatasetDependencyNode] = DatasetDependencyNode
     job_class: type[Job] = Job
+    checkpoint_class: type[Checkpoint] = Checkpoint
 
     def __init__(
         self,
-        uri: Optional[StorageURI] = None,
+        uri: StorageURI | None = None,
     ):
         self.uri = uri or StorageURI("")
 
@@ -89,7 +108,7 @@ class AbstractMetastore(ABC, Serializable):
     @abstractmethod
     def clone(
         self,
-        uri: Optional[StorageURI] = None,
+        uri: StorageURI | None = None,
         use_new_connection: bool = False,
     ) -> "AbstractMetastore":
         """Clones AbstractMetastore implementation for some Storage input.
@@ -106,6 +125,16 @@ class AbstractMetastore(ABC, Serializable):
         differently."""
         self.close()
 
+    @contextmanager
+    def _init_guard(self):
+        """Ensure resources acquired during __init__ are released on failure."""
+        try:
+            yield
+        except Exception:
+            with suppress(Exception):
+                self.close_on_exit()
+            raise
+
     def cleanup_tables(self, temp_table_names: list[str]) -> None:
         """Cleanup temp tables."""
 
@@ -129,8 +158,8 @@ class AbstractMetastore(ABC, Serializable):
     def create_namespace(
         self,
         name: str,
-        description: Optional[str] = None,
-        uuid: Optional[str] = None,
+        description: str | None = None,
+        uuid: str | None = None,
         ignore_if_exists: bool = True,
         validate: bool = True,
         **kwargs,
@@ -141,6 +170,10 @@ class AbstractMetastore(ABC, Serializable):
     def get_namespace(self, name: str, conn=None) -> Namespace:
         """Gets a single namespace by name"""
 
+    @abstractmethod
+    def remove_namespace(self, namespace_id: int, conn=None) -> None:
+        """Removes a single namespace by id"""
+
     @abstractmethod
     def list_namespaces(self, conn=None) -> list[Namespace]:
         """Gets a list of all namespaces"""
@@ -173,8 +206,8 @@ class AbstractMetastore(ABC, Serializable):
         self,
         namespace_name: str,
         name: str,
-        description: Optional[str] = None,
-        uuid: Optional[str] = None,
+        description: str | None = None,
+        uuid: str | None = None,
         ignore_if_exists: bool = True,
         validate: bool = True,
         **kwargs,
@@ -190,12 +223,32 @@ class AbstractMetastore(ABC, Serializable):
         It also creates project if not found and create flag is set to True.
         """
 
+    def is_default_project(self, project_name: str, namespace_name: str) -> bool:
+        return (
+            project_name == self.default_project_name
+            and namespace_name == self.default_namespace_name
+        )
+
+    def is_listing_project(self, project_name: str, namespace_name: str) -> bool:
+        return (
+            project_name == self.listing_project_name
+            and namespace_name == self.system_namespace_name
+        )
+
     @abstractmethod
     def get_project_by_id(self, project_id: int, conn=None) -> Project:
         """Gets a single project by id"""
 
     @abstractmethod
-    def list_projects(self, namespace_id: Optional[int], conn=None) -> list[Project]:
+    def count_projects(self, namespace_id: int | None = None) -> int:
+        """Counts projects in some namespace or in general."""
+
+    @abstractmethod
+    def remove_project(self, project_id: int, conn=None) -> None:
+        """Removes a single project by id"""
+
+    @abstractmethod
+    def list_projects(self, namespace_id: int | None, conn=None) -> list[Project]:
         """Gets list of projects in some namespace or in general (in all namespaces)"""
 
     #
@@ -205,15 +258,15 @@ class AbstractMetastore(ABC, Serializable):
     def create_dataset(
         self,
         name: str,
-        project_id: Optional[int] = None,
+        project_id: int | None = None,
         status: int = DatasetStatus.CREATED,
-        sources: Optional[list[str]] = None,
-        feature_schema: Optional[dict] = None,
+        sources: list[str] | None = None,
+        feature_schema: dict | None = None,
         query_script: str = "",
-        schema: Optional[dict[str, Any]] = None,
+        schema: dict[str, Any] | None = None,
         ignore_if_exists: bool = False,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
     ) -> DatasetRecord:
         """Creates new dataset."""
 
@@ -224,20 +277,20 @@ class AbstractMetastore(ABC, Serializable):
         version: str,
         status: int,
         sources: str = "",
-        feature_schema: Optional[dict] = None,
+        feature_schema: dict | None = None,
         query_script: str = "",
         error_message: str = "",
         error_stack: str = "",
         script_output: str = "",
-        created_at: Optional[datetime] = None,
-        finished_at: Optional[datetime] = None,
-        schema: Optional[dict[str, Any]] = None,
+        created_at: datetime | None = None,
+        finished_at: datetime | None = None,
+        schema: dict[str, Any] | None = None,
         ignore_if_exists: bool = False,
-        num_objects: Optional[int] = None,
-        size: Optional[int] = None,
-        preview: Optional[list[dict]] = None,
-        job_id: Optional[str] = None,
-        uuid: Optional[str] = None,
+        num_objects: int | None = None,
+        size: int | None = None,
+        preview: list[dict] | None = None,
+        job_id: str | None = None,
+        uuid: str | None = None,
     ) -> DatasetRecord:
         """Creates new dataset version."""
 
@@ -266,13 +319,17 @@ class AbstractMetastore(ABC, Serializable):
 
     @abstractmethod
     def list_datasets(
-        self, project_id: Optional[int] = None
+        self, project_id: int | None = None
     ) -> Iterator[DatasetListRecord]:
         """Lists all datasets in some project or in all projects."""
 
+    @abstractmethod
+    def count_datasets(self, project_id: int | None = None) -> int:
+        """Counts datasets in some project or in all projects."""
+
     @abstractmethod
     def list_datasets_by_prefix(
-        self, prefix: str, project_id: Optional[int] = None
+        self, prefix: str, project_id: int | None = None
     ) -> Iterator["DatasetListRecord"]:
         """
         Lists all datasets which names start with prefix in some project or in all
@@ -283,8 +340,8 @@ class AbstractMetastore(ABC, Serializable):
     def get_dataset(
         self,
         name: str,  # normal, not full dataset name
-        namespace_name: Optional[str] = None,
-        project_name: Optional[str] = None,
+        namespace_name: str | None = None,
+        project_name: str | None = None,
         conn=None,
     ) -> DatasetRecord:
         """Gets a single dataset by name."""
@@ -294,7 +351,7 @@ class AbstractMetastore(ABC, Serializable):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[str] = None,
+        version: str | None = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -319,20 +376,26 @@ class AbstractMetastore(ABC, Serializable):
         self,
         source_dataset: DatasetRecord,
         source_dataset_version: str,
-        new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[str] = None,
+        new_source_dataset: DatasetRecord | None = None,
+        new_source_dataset_version: str | None = None,
     ) -> None:
         """Updates dataset dependency source."""
 
     @abstractmethod
     def get_direct_dataset_dependencies(
         self, dataset: DatasetRecord, version: str
-    ) -> list[Optional[DatasetDependency]]:
+    ) -> list[DatasetDependency | None]:
         """Gets direct dataset dependencies."""
 
+    @abstractmethod
+    def get_dataset_dependency_nodes(
+        self, dataset_id: int, version_id: int
+    ) -> list[DatasetDependencyNode | None]:
+        """Gets dataset dependency node from database."""
+
     @abstractmethod
     def remove_dataset_dependencies(
-        self, dataset: DatasetRecord, version: Optional[str] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well.
@@ -340,7 +403,7 @@ class AbstractMetastore(ABC, Serializable):
 
     @abstractmethod
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[str] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -362,8 +425,9 @@ class AbstractMetastore(ABC, Serializable):
         query_type: JobQueryType = JobQueryType.PYTHON,
         status: JobStatus = JobStatus.CREATED,
         workers: int = 1,
-        python_version: Optional[str] = None,
-        params: Optional[dict[str, str]] = None,
+        python_version: str | None = None,
+        params: dict[str, str] | None = None,
+        parent_job_id: str | None = None,
     ) -> str:
         """
         Creates a new job.
@@ -371,19 +435,19 @@ class AbstractMetastore(ABC, Serializable):
         """
 
     @abstractmethod
-    def get_job(self, job_id: str) -> Optional[Job]:
+    def get_job(self, job_id: str) -> Job | None:
         """Returns the job with the given ID."""
 
     @abstractmethod
     def update_job(
         self,
         job_id: str,
-        status: Optional[JobStatus] = None,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
-        finished_at: Optional[datetime] = None,
-        metrics: Optional[dict[str, Any]] = None,
-    ) -> Optional["Job"]:
+        status: JobStatus | None = None,
+        error_message: str | None = None,
+        error_stack: str | None = None,
+        finished_at: datetime | None = None,
+        metrics: dict[str, Any] | None = None,
+    ) -> Job | None:
         """Updates job fields."""
 
     @abstractmethod
@@ -391,15 +455,90 @@ class AbstractMetastore(ABC, Serializable):
         self,
         job_id: str,
         status: JobStatus,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
+        error_message: str | None = None,
+        error_stack: str | None = None,
     ) -> None:
         """Set the status of the given job."""
 
     @abstractmethod
-    def get_job_status(self, job_id: str) -> Optional[JobStatus]:
+    def get_job_status(self, job_id: str) -> JobStatus | None:
         """Returns the status of the given job."""
 
+    @abstractmethod
+    def get_last_job_by_name(self, name: str, conn=None) -> "Job | None":
+        """Returns the last job with the given name, ordered by created_at."""
+
+    #
+    # Checkpoints
+    #
+
+    @abstractmethod
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
+        """Returns all checkpoints related to some job"""
+
+    @abstractmethod
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Checkpoint | None:
+        """Get last created checkpoint for some job."""
+
+    @abstractmethod
+    def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
+        """Gets single checkpoint by id"""
+
+    def find_checkpoint(
+        self, job_id: str, _hash: str, partial: bool = False, conn=None
+    ) -> Checkpoint | None:
+        """
+        Tries to find checkpoint for a job with specific hash and optionally partial
+        """
+
+    @abstractmethod
+    def create_checkpoint(
+        self,
+        job_id: str,
+        _hash: str,
+        partial: bool = False,
+        conn: Any | None = None,
+    ) -> Checkpoint:
+        """Creates new checkpoint"""
+
+    #
+    # Dataset Version Jobs (many-to-many)
+    #
+
+    @abstractmethod
+    def link_dataset_version_to_job(
+        self,
+        dataset_version_id: int,
+        job_id: str,
+        is_creator: bool = False,
+        conn=None,
+    ) -> None:
+        """
+        Link dataset version to job.
+
+        This atomically:
+        1. Creates a link in the dataset_version_jobs junction table
+        2. Updates dataset_version.job_id to point to this job
+        """
+
+    @abstractmethod
+    def get_ancestor_job_ids(self, job_id: str, conn=None) -> list[str]:
+        """Get all ancestor job IDs for a given job."""
+
+    @abstractmethod
+    def get_dataset_version_for_job_ancestry(
+        self,
+        dataset_name: str,
+        namespace_name: str,
+        project_name: str,
+        job_id: str,
+        conn=None,
+    ) -> DatasetVersion | None:
+        """
+        Find the dataset version that was created by any job in the ancestry.
+        Returns the most recently linked version from these jobs.
+        """
+
 
 class AbstractDBMetastore(AbstractMetastore):
     """
@@ -414,11 +553,13 @@ class AbstractDBMetastore(AbstractMetastore):
     DATASET_TABLE = "datasets"
     DATASET_VERSION_TABLE = "datasets_versions"
     DATASET_DEPENDENCY_TABLE = "datasets_dependencies"
+    DATASET_VERSION_JOBS_TABLE = "dataset_version_jobs"
     JOBS_TABLE = "jobs"
+    CHECKPOINTS_TABLE = "checkpoints"
 
     db: "DatabaseEngine"
 
-    def __init__(self, uri: Optional[StorageURI] = None):
+    def __init__(self, uri: StorageURI | None = None):
         uri = uri or StorageURI("")
         super().__init__(uri)
 
@@ -658,9 +799,6 @@ class AbstractDBMetastore(AbstractMetastore):
             return self._projects.select()
         return select(*columns)
 
-    def _projects_update(self) -> "Update":
-        return self._projects.update()
-
     def _projects_delete(self) -> "Delete":
         return self._projects.delete()
 
@@ -713,8 +851,8 @@ class AbstractDBMetastore(AbstractMetastore):
     def create_namespace(
         self,
         name: str,
-        description: Optional[str] = None,
-        uuid: Optional[str] = None,
+        description: str | None = None,
+        uuid: str | None = None,
         ignore_if_exists: bool = True,
         validate: bool = True,
         **kwargs,
@@ -735,6 +873,18 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.get_namespace(name)
 
+    def remove_namespace(self, namespace_id: int, conn=None) -> None:
+        num_projects = self.count_projects(namespace_id)
+        if num_projects > 0:
+            raise NamespaceDeleteNotAllowedError(
+                f"Namespace cannot be removed. It contains {num_projects} project(s). "
+                "Please remove the project(s) first."
+            )
+
+        n = self._namespaces
+        with self.db.transaction():
+            self.db.execute(self._namespaces_delete().where(n.c.id == namespace_id))
+
     def get_namespace(self, name: str, conn=None) -> Namespace:
         """Gets a single namespace by name"""
         n = self._namespaces
@@ -766,8 +916,8 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         namespace_name: str,
         name: str,
-        description: Optional[str] = None,
-        uuid: Optional[str] = None,
+        description: str | None = None,
+        uuid: str | None = None,
         ignore_if_exists: bool = True,
         validate: bool = True,
         **kwargs,
@@ -796,17 +946,15 @@ class AbstractDBMetastore(AbstractMetastore):
 
         return self.get_project(name, namespace.name)
 
-    def _is_listing_project(self, project_name: str, namespace_name: str) -> bool:
-        return (
-            project_name == self.listing_project_name
-            and namespace_name == self.system_namespace_name
-        )
+    def _projects_base_query(self) -> "Select":
+        n = self._namespaces
+        p = self._projects
 
-    def _is_default_project(self, project_name: str, namespace_name: str) -> bool:
-        return (
-            project_name == self.default_project_name
-            and namespace_name == self.default_namespace_name
+        query = self._projects_select(
+            *(getattr(n.c, f) for f in self._namespaces_fields),
+            *(getattr(p.c, f) for f in self._projects_fields),
         )
+        return query.select_from(n.join(p, n.c.id == p.c.namespace_id))
 
     def get_project(
         self, name: str, namespace_name: str, create: bool = False, conn=None
@@ -816,18 +964,14 @@ class AbstractDBMetastore(AbstractMetastore):
         p = self._projects
         validate = True
 
-        if self._is_listing_project(name, namespace_name) or self._is_default_project(
+        if self.is_listing_project(name, namespace_name) or self.is_default_project(
             name, namespace_name
         ):
             # we are always creating default and listing projects if they don't exist
             create = True
             validate = False
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id)).where(
+        query = self._projects_base_query().where(
             p.c.name == name, n.c.name == namespace_name
         )
 
@@ -842,37 +986,50 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def get_project_by_id(self, project_id: int, conn=None) -> Project:
         """Gets a single project by id"""
-        n = self._namespaces
         p = self._projects
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id)).where(
-            p.c.id == project_id
-        )
+        query = self._projects_base_query().where(p.c.id == project_id)
 
         rows = list(self.db.execute(query, conn=conn))
         if not rows:
             raise ProjectNotFoundError(f"Project with id {project_id} not found.")
         return self.project_class.parse(*rows[0])
 
-    def list_projects(self, namespace_id: Optional[int], conn=None) -> list[Project]:
+    def count_projects(self, namespace_id: int | None = None) -> int:
+        p = self._projects
+
+        query = self._projects_base_query()
+        if namespace_id:
+            query = query.where(p.c.namespace_id == namespace_id)
+
+        query = select(f.count(1)).select_from(query.subquery())
+
+        return next(self.db.execute(query))[0]
+
+    def remove_project(self, project_id: int, conn=None) -> None:
+        num_datasets = self.count_datasets(project_id)
+        if num_datasets > 0:
+            raise ProjectDeleteNotAllowedError(
+                f"Project cannot be removed. It contains {num_datasets} dataset(s). "
+                "Please remove the dataset(s) first."
+            )
+
+        p = self._projects
+        with self.db.transaction():
+            self.db.execute(self._projects_delete().where(p.c.id == project_id))
+
+    def list_projects(
+        self, namespace_id: int | None = None, conn=None
+    ) -> list[Project]:
         """
         Gets a list of projects inside some namespace, or in all namespaces
         """
-        n = self._namespaces
         p = self._projects
 
-        query = self._projects_select(
-            *(getattr(n.c, f) for f in self._namespaces_fields),
-            *(getattr(p.c, f) for f in self._projects_fields),
-        )
-        query = query.select_from(n.join(p, n.c.id == p.c.namespace_id))
+        query = self._projects_base_query()
 
         if namespace_id:
-            query = query.where(n.c.id == namespace_id)
+            query = query.where(p.c.namespace_id == namespace_id)
 
         rows = list(self.db.execute(query, conn=conn))
 
@@ -885,15 +1042,15 @@ class AbstractDBMetastore(AbstractMetastore):
     def create_dataset(
         self,
         name: str,
-        project_id: Optional[int] = None,
+        project_id: int | None = None,
         status: int = DatasetStatus.CREATED,
-        sources: Optional[list[str]] = None,
-        feature_schema: Optional[dict] = None,
+        sources: list[str] | None = None,
+        feature_schema: dict | None = None,
         query_script: str = "",
-        schema: Optional[dict[str, Any]] = None,
+        schema: dict[str, Any] | None = None,
         ignore_if_exists: bool = False,
-        description: Optional[str] = None,
-        attrs: Optional[list[str]] = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
         **kwargs,  # TODO registered = True / False
     ) -> DatasetRecord:
         """Creates new dataset."""
@@ -933,20 +1090,20 @@ class AbstractDBMetastore(AbstractMetastore):
         version: str,
         status: int,
         sources: str = "",
-        feature_schema: Optional[dict] = None,
+        feature_schema: dict | None = None,
         query_script: str = "",
         error_message: str = "",
         error_stack: str = "",
         script_output: str = "",
-        created_at: Optional[datetime] = None,
-        finished_at: Optional[datetime] = None,
-        schema: Optional[dict[str, Any]] = None,
+        created_at: datetime | None = None,
+        finished_at: datetime | None = None,
+        schema: dict[str, Any] | None = None,
         ignore_if_exists: bool = False,
-        num_objects: Optional[int] = None,
-        size: Optional[int] = None,
-        preview: Optional[list[dict]] = None,
-        job_id: Optional[str] = None,
-        uuid: Optional[str] = None,
+        num_objects: int | None = None,
+        size: int | None = None,
+        preview: list[dict] | None = None,
+        job_id: str | None = None,
+        uuid: str | None = None,
         conn=None,
     ) -> DatasetRecord:
         """Creates new dataset version."""
@@ -1024,7 +1181,7 @@ class AbstractDBMetastore(AbstractMetastore):
                     dataset_values[field] = None
                 else:
                     values[field] = json.dumps(value)
-                    dataset_values[field] = DatasetRecord.parse_schema(value)
+                    dataset_values[field] = parse_schema(value)
             elif field == "project_id":
                 if not value:
                     raise ValueError("Cannot set empty project_id for dataset")
@@ -1075,9 +1232,7 @@ class AbstractDBMetastore(AbstractMetastore):
 
             if field == "schema":
                 values[field] = json.dumps(value) if value else None
-                version_values[field] = (
-                    DatasetRecord.parse_schema(value) if value else None
-                )
+                version_values[field] = parse_schema(value) if value else None
             elif field == "feature_schema":
                 if value is None:
                     values[field] = None
@@ -1092,7 +1247,7 @@ class AbstractDBMetastore(AbstractMetastore):
                         f"Field '{field}' must be a list, got {type(value).__name__}"
                     )
                 else:
-                    values[field] = json.dumps(value, cls=JSONSerialize)
+                    values[field] = json.dumps(value, serialize_bytes=True)
                 version_values["_preview_data"] = value
             else:
                 values[field] = value
@@ -1118,13 +1273,13 @@ class AbstractDBMetastore(AbstractMetastore):
             f"Dataset {dataset.name} does not have version {version}"
         )
 
-    def _parse_dataset(self, rows) -> Optional[DatasetRecord]:
+    def _parse_dataset(self, rows) -> DatasetRecord | None:
         versions = [self.dataset_class.parse(*r) for r in rows]
         if not versions:
             return None
         return reduce(lambda ds, version: ds.merge_versions(version), versions)
 
-    def _parse_list_dataset(self, rows) -> Optional[DatasetListRecord]:
+    def _parse_list_dataset(self, rows) -> DatasetListRecord | None:
         versions = [self.dataset_list_class.parse(*r) for r in rows]
         if not versions:
             return None
@@ -1187,9 +1342,8 @@ class AbstractDBMetastore(AbstractMetastore):
         )
 
     def list_datasets(
-        self, project_id: Optional[int] = None
+        self, project_id: int | None = None
     ) -> Iterator["DatasetListRecord"]:
-        """Lists all datasets."""
         d = self._datasets
         query = self._base_list_datasets_query().order_by(
             self._datasets.c.name, self._datasets_versions.c.version
@@ -1198,8 +1352,18 @@ class AbstractDBMetastore(AbstractMetastore):
             query = query.where(d.c.project_id == project_id)
         yield from self._parse_dataset_list(self.db.execute(query))
 
+    def count_datasets(self, project_id: int | None = None) -> int:
+        d = self._datasets
+        query = self._datasets_select()
+        if project_id:
+            query = query.where(d.c.project_id == project_id)
+
+        query = select(f.count(1)).select_from(query.subquery())
+
+        return next(self.db.execute(query))[0]
+
     def list_datasets_by_prefix(
-        self, prefix: str, project_id: Optional[int] = None, conn=None
+        self, prefix: str, project_id: int | None = None, conn=None
     ) -> Iterator["DatasetListRecord"]:
         d = self._datasets
         query = self._base_list_datasets_query()
@@ -1211,8 +1375,8 @@ class AbstractDBMetastore(AbstractMetastore):
     def get_dataset(
         self,
         name: str,  # normal, not full dataset name
-        namespace_name: Optional[str] = None,
-        project_name: Optional[str] = None,
+        namespace_name: str | None = None,
+        project_name: str | None = None,
         conn=None,
     ) -> DatasetRecord:
         """
@@ -1273,7 +1437,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[str] = None,
+        version: str | None = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -1327,8 +1491,8 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         source_dataset: DatasetRecord,
         source_dataset_version: str,
-        new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[str] = None,
+        new_source_dataset: DatasetRecord | None = None,
+        new_source_dataset_version: str | None = None,
     ) -> None:
         dd = self._datasets_dependencies
 
@@ -1358,9 +1522,21 @@ class AbstractDBMetastore(AbstractMetastore):
         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[Optional[DatasetDependency]]:
+    ) -> list[DatasetDependency | None]:
         n = self._namespaces
         p = self._projects
         d = self._datasets
@@ -1387,8 +1563,77 @@ class AbstractDBMetastore(AbstractMetastore):
 
         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, depth_limit: int = DEPTH_LIMIT_DEFAULT
+    ) -> list[DatasetDependencyNode | None]:
+        n = self._namespaces_select().subquery()
+        p = self._projects
+        d = self._datasets_select().subquery()
+        dd = self._datasets_dependencies
+        dv = self._datasets_versions
+
+        # Common dependency fields for CTE
+        dep_fields = [
+            dd.c.id,
+            dd.c.source_dataset_id,
+            dd.c.source_dataset_version_id,
+            dd.c.dataset_id,
+            dd.c.dataset_version_id,
+        ]
+
+        # Base case: direct dependencies
+        base_query = select(
+            *dep_fields,
+            literal(0).label("depth"),
+        ).where(
+            (dd.c.source_dataset_id == dataset_id)
+            & (dd.c.source_dataset_version_id == version_id)
+        )
+
+        cte = base_query.cte(name="dependency_tree", recursive=True)
+
+        # Recursive case: dependencies of dependencies
+        # 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
+        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
+            # dependencies visible to callers.
+            cte.join(d, cte.c.dataset_id == d.c.id, isouter=True)
+            .join(dv, cte.c.dataset_version_id == dv.c.id, isouter=True)
+            .join(p, d.c.project_id == p.c.id, isouter=True)
+            .join(n, p.c.namespace_id == n.c.id, isouter=True)
+        )
+
+        return [
+            self.dependency_node_class.parse(*r) for r in self.db.execute(final_query)
+        ]
+
     def remove_dataset_dependencies(
-        self, dataset: DatasetRecord, version: Optional[str] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well
@@ -1407,7 +1652,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self.db.execute(q)
 
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[str] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -1458,11 +1703,13 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("error_stack", Text, nullable=False, default=""),
             Column("params", JSON, nullable=False),
             Column("metrics", JSON, nullable=False),
+            Column("parent_job_id", Text, nullable=True),
+            Index("idx_jobs_parent_job_id", "parent_job_id"),
         ]
 
     @cached_property
     def _job_fields(self) -> list[str]:
-        return [c.name for c in self._jobs_columns() if c.name]  # type: ignore[attr-defined]
+        return [c.name for c in self._jobs_columns() if isinstance(c, Column)]  # type: ignore[attr-defined]
 
     @cached_property
     def _jobs(self) -> "Table":
@@ -1496,6 +1743,18 @@ class AbstractDBMetastore(AbstractMetastore):
         query = self._jobs_query().where(self._jobs.c.id.in_(ids))
         yield from self._parse_jobs(self.db.execute(query, conn=conn))
 
+    def get_last_job_by_name(self, name: str, conn=None) -> "Job | None":
+        query = (
+            self._jobs_query()
+            .where(self._jobs.c.name == name)
+            .order_by(self._jobs.c.created_at.desc())
+            .limit(1)
+        )
+        results = list(self.db.execute(query, conn=conn))
+        if not results:
+            return None
+        return self._parse_job(results[0])
+
     def create_job(
         self,
         name: str,
@@ -1503,9 +1762,10 @@ class AbstractDBMetastore(AbstractMetastore):
         query_type: JobQueryType = JobQueryType.PYTHON,
         status: JobStatus = JobStatus.CREATED,
         workers: int = 1,
-        python_version: Optional[str] = None,
-        params: Optional[dict[str, str]] = None,
-        conn: Optional[Any] = None,
+        python_version: str | None = None,
+        params: dict[str, str] | None = None,
+        parent_job_id: str | None = None,
+        conn: Any = None,
     ) -> str:
         """
         Creates a new job.
@@ -1526,12 +1786,13 @@ class AbstractDBMetastore(AbstractMetastore):
                 error_stack="",
                 params=json.dumps(params or {}),
                 metrics=json.dumps({}),
+                parent_job_id=parent_job_id,
             ),
             conn=conn,
         )
         return job_id
 
-    def get_job(self, job_id: str, conn=None) -> Optional[Job]:
+    def get_job(self, job_id: str, conn=None) -> Job | None:
         """Returns the job with the given ID."""
         query = self._jobs_select(self._jobs).where(self._jobs.c.id == job_id)
         results = list(self.db.execute(query, conn=conn))
@@ -1542,13 +1803,13 @@ class AbstractDBMetastore(AbstractMetastore):
     def update_job(
         self,
         job_id: str,
-        status: Optional[JobStatus] = None,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
-        finished_at: Optional[datetime] = None,
-        metrics: Optional[dict[str, Any]] = None,
-        conn: Optional[Any] = None,
-    ) -> Optional["Job"]:
+        status: JobStatus | None = None,
+        error_message: str | None = None,
+        error_stack: str | None = None,
+        finished_at: datetime | None = None,
+        metrics: dict[str, Any] | None = None,
+        conn: Any | None = None,
+    ) -> Job | None:
         """Updates job fields."""
         values: dict = {}
         if status is not None:
@@ -1575,9 +1836,9 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         job_id: str,
         status: JobStatus,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
-        conn: Optional[Any] = None,
+        error_message: str | None = None,
+        error_stack: str | None = None,
+        conn: Any | None = None,
     ) -> None:
         """Set the status of the given job."""
         values: dict = {"status": status}
@@ -1595,8 +1856,8 @@ class AbstractDBMetastore(AbstractMetastore):
     def get_job_status(
         self,
         job_id: str,
-        conn: Optional[Any] = None,
-    ) -> Optional[JobStatus]:
+        conn: Any | None = None,
+    ) -> JobStatus | None:
         """Returns the status of the given job."""
         results = list(
             self.db.execute(
@@ -1607,3 +1868,321 @@ class AbstractDBMetastore(AbstractMetastore):
         if not results:
             return None
         return results[0][0]
+
+    #
+    # Checkpoints
+    #
+
+    @staticmethod
+    def _checkpoints_columns() -> "list[SchemaItem]":
+        return [
+            Column(
+                "id",
+                Text,
+                default=uuid4,
+                primary_key=True,
+                nullable=False,
+            ),
+            Column("job_id", Text, nullable=True),
+            Column("hash", Text, nullable=False),
+            Column("partial", Boolean, default=False),
+            Column("created_at", DateTime(timezone=True), nullable=False),
+            UniqueConstraint("job_id", "hash"),
+        ]
+
+    @cached_property
+    def _checkpoints_fields(self) -> list[str]:
+        return [c.name for c in self._checkpoints_columns() if c.name]  # type: ignore[attr-defined]
+
+    @cached_property
+    def _checkpoints(self) -> "Table":
+        return Table(
+            self.CHECKPOINTS_TABLE,
+            self.db.metadata,
+            *self._checkpoints_columns(),
+        )
+
+    @abstractmethod
+    def _checkpoints_insert(self) -> "Insert": ...
+
+    @classmethod
+    def _dataset_version_jobs_columns(cls) -> "list[SchemaItem]":
+        """Junction table for dataset versions and jobs many-to-many relationship."""
+        return [
+            Column("id", Integer, primary_key=True),
+            Column(
+                "dataset_version_id",
+                Integer,
+                ForeignKey(f"{cls.DATASET_VERSION_TABLE}.id", ondelete="CASCADE"),
+                nullable=False,
+            ),
+            Column("job_id", Text, nullable=False),
+            Column("is_creator", Boolean, nullable=False, default=False),
+            Column("created_at", DateTime(timezone=True)),
+            UniqueConstraint("dataset_version_id", "job_id"),
+            Index("dc_idx_dvj_query", "job_id", "is_creator", "created_at"),
+        ]
+
+    @cached_property
+    def _dataset_version_jobs_fields(self) -> list[str]:
+        return [c.name for c in self._dataset_version_jobs_columns() if c.name]  # type: ignore[attr-defined]
+
+    @cached_property
+    def _dataset_version_jobs(self) -> "Table":
+        return Table(
+            self.DATASET_VERSION_JOBS_TABLE,
+            self.db.metadata,
+            *self._dataset_version_jobs_columns(),
+        )
+
+    @abstractmethod
+    def _dataset_version_jobs_insert(self) -> "Insert": ...
+
+    def _dataset_version_jobs_select(self, *columns) -> "Select":
+        if not columns:
+            return self._dataset_version_jobs.select()
+        return select(*columns)
+
+    def _dataset_version_jobs_delete(self) -> "Delete":
+        return self._dataset_version_jobs.delete()
+
+    def _checkpoints_select(self, *columns) -> "Select":
+        if not columns:
+            return self._checkpoints.select()
+        return select(*columns)
+
+    def _checkpoints_delete(self) -> "Delete":
+        return self._checkpoints.delete()
+
+    def _checkpoints_query(self):
+        return self._checkpoints_select(
+            *[getattr(self._checkpoints.c, f) for f in self._checkpoints_fields]
+        )
+
+    def create_checkpoint(
+        self,
+        job_id: str,
+        _hash: str,
+        partial: bool = False,
+        conn: Any | None = None,
+    ) -> Checkpoint:
+        """
+        Creates a new job query step.
+        """
+        checkpoint_id = str(uuid4())
+        self.db.execute(
+            self._checkpoints_insert().values(
+                id=checkpoint_id,
+                job_id=job_id,
+                hash=_hash,
+                partial=partial,
+                created_at=datetime.now(timezone.utc),
+            ),
+            conn=conn,
+        )
+        return self.get_checkpoint_by_id(checkpoint_id)
+
+    def list_checkpoints(self, job_id: str, conn=None) -> Iterator[Checkpoint]:
+        """List checkpoints by job id."""
+        query = self._checkpoints_query().where(self._checkpoints.c.job_id == job_id)
+        rows = list(self.db.execute(query, conn=conn))
+
+        yield from [self.checkpoint_class.parse(*r) for r in rows]
+
+    def get_checkpoint_by_id(self, checkpoint_id: str, conn=None) -> Checkpoint:
+        """Returns the checkpoint with the given ID."""
+        ch = self._checkpoints
+        query = self._checkpoints_select(ch).where(ch.c.id == checkpoint_id)
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            raise CheckpointNotFoundError(f"Checkpoint {checkpoint_id} not found")
+        return self.checkpoint_class.parse(*rows[0])
+
+    def find_checkpoint(
+        self, job_id: str, _hash: str, partial: bool = False, conn=None
+    ) -> Checkpoint | None:
+        """
+        Tries to find checkpoint for a job with specific hash and optionally partial
+        """
+        ch = self._checkpoints
+        query = self._checkpoints_select(ch).where(
+            ch.c.job_id == job_id, ch.c.hash == _hash, ch.c.partial == partial
+        )
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            return None
+        return self.checkpoint_class.parse(*rows[0])
+
+    def get_last_checkpoint(self, job_id: str, conn=None) -> Checkpoint | None:
+        query = (
+            self._checkpoints_query()
+            .where(self._checkpoints.c.job_id == job_id)
+            .order_by(desc(self._checkpoints.c.created_at))
+            .limit(1)
+        )
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            return None
+        return self.checkpoint_class.parse(*rows[0])
+
+    def link_dataset_version_to_job(
+        self,
+        dataset_version_id: int,
+        job_id: str,
+        is_creator: bool = False,
+        conn=None,
+    ) -> None:
+        # Use transaction to atomically:
+        # 1. Link dataset version to job in junction table
+        # 2. Update dataset_version.job_id to point to this job
+        with self.db.transaction() as tx_conn:
+            conn = conn or tx_conn
+
+            # Insert into junction table
+            query = self._dataset_version_jobs_insert().values(
+                dataset_version_id=dataset_version_id,
+                job_id=job_id,
+                is_creator=is_creator,
+                created_at=datetime.now(timezone.utc),
+            )
+            if hasattr(query, "on_conflict_do_nothing"):
+                query = query.on_conflict_do_nothing(
+                    index_elements=["dataset_version_id", "job_id"]
+                )
+            self.db.execute(query, conn=conn)
+
+            # Also update dataset_version.job_id to point to this job
+            update_query = (
+                self._datasets_versions.update()
+                .where(self._datasets_versions.c.id == dataset_version_id)
+                .values(job_id=job_id)
+            )
+            self.db.execute(update_query, conn=conn)
+
+    def get_ancestor_job_ids(self, job_id: str, conn=None) -> list[str]:
+        # Use recursive CTE to walk up the parent chain
+        # Format: WITH RECURSIVE ancestors(id, parent_job_id, depth) AS (...)
+        # Include depth tracking to prevent infinite recursion in case of
+        # circular dependencies
+        ancestors_cte = (
+            self._jobs_select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+                literal(0).label("depth"),
+            )
+            .where(self._jobs.c.id == job_id)
+            .cte(name="ancestors", recursive=True)
+        )
+
+        # Recursive part: join with parent jobs, incrementing depth and checking limit
+        ancestors_recursive = ancestors_cte.union_all(
+            self._jobs_select(
+                self._jobs.c.id.label("id"),
+                self._jobs.c.parent_job_id.label("parent_job_id"),
+                (ancestors_cte.c.depth + 1).label("depth"),
+            ).select_from(
+                self._jobs.join(
+                    ancestors_cte,
+                    (
+                        self._jobs.c.id
+                        == cast(ancestors_cte.c.parent_job_id, self._jobs.c.id.type)
+                    )
+                    & (ancestors_cte.c.parent_job_id.isnot(None))  # Stop at root jobs
+                    & (ancestors_cte.c.depth < JOB_ANCESTRY_MAX_DEPTH),
+                )
+            )
+        )
+
+        # Select all ancestor IDs and depths except the starting job itself
+        query = select(ancestors_recursive.c.id, ancestors_recursive.c.depth).where(
+            ancestors_recursive.c.id != job_id
+        )
+
+        results = list(self.db.execute(query, conn=conn))
+
+        # Check if we hit the depth limit
+        if results:
+            max_found_depth = max(row[1] for row in results)
+            if max_found_depth >= JOB_ANCESTRY_MAX_DEPTH:
+                from datachain.error import JobAncestryDepthExceededError
+
+                raise JobAncestryDepthExceededError(
+                    f"Job ancestry chain exceeds maximum depth of "
+                    f"{JOB_ANCESTRY_MAX_DEPTH}. Job ID: {job_id}"
+                )
+
+        return [str(row[0]) for row in results]
+
+    def _get_dataset_version_for_job_ancestry_query(
+        self,
+        dataset_name: str,
+        namespace_name: str,
+        project_name: str,
+        job_ancestry: list[str],
+    ) -> "Select":
+        """Find most recent dataset version created by any job in ancestry.
+
+        Searches job ancestry (current + parents) for the newest version of
+        the dataset where is_creator=True. Returns newest by created_at, or
+        None if no version was created by any job in the ancestry chain.
+
+        Used for checkpoint resolution to find which version to reuse when
+        continuing from a parent job.
+        """
+        return (
+            self._datasets_versions_select()
+            .select_from(
+                self._dataset_version_jobs.join(
+                    self._datasets_versions,
+                    self._dataset_version_jobs.c.dataset_version_id
+                    == self._datasets_versions.c.id,
+                )
+                .join(
+                    self._datasets,
+                    self._datasets_versions.c.dataset_id == self._datasets.c.id,
+                )
+                .join(
+                    self._projects,
+                    self._datasets.c.project_id == self._projects.c.id,
+                )
+                .join(
+                    self._namespaces,
+                    self._projects.c.namespace_id == self._namespaces.c.id,
+                )
+            )
+            .where(
+                self._datasets.c.name == dataset_name,
+                self._namespaces.c.name == namespace_name,
+                self._projects.c.name == project_name,
+                self._dataset_version_jobs.c.job_id.in_(job_ancestry),
+                self._dataset_version_jobs.c.is_creator.is_(True),
+            )
+            .order_by(desc(self._dataset_version_jobs.c.created_at))
+            .limit(1)
+        )
+
+    def get_dataset_version_for_job_ancestry(
+        self,
+        dataset_name: str,
+        namespace_name: str,
+        project_name: str,
+        job_id: str,
+        conn=None,
+    ) -> DatasetVersion | None:
+        # Get job ancestry (current job + all ancestors)
+        job_ancestry = [job_id, *self.get_ancestor_job_ids(job_id, conn=conn)]
+
+        query = self._get_dataset_version_for_job_ancestry_query(
+            dataset_name, namespace_name, project_name, job_ancestry
+        )
+
+        results = list(self.db.execute(query, conn=conn))
+        if not results:
+            return None
+
+        if len(results) > 1:
+            raise DataChainError(
+                f"Expected at most 1 dataset version, found {len(results)}"
+            )
+
+        return self.dataset_version_class.parse(*results[0])