datachain 0.14.2__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,22 +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.utils import JSONSerialize
+from datachain.namespace import Namespace
+from datachain.project import Project
 
 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):
@@ -60,15 +81,20 @@ class AbstractMetastore(ABC, Serializable):
     uri: StorageURI
 
     schema: "schema.Schema"
+    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("")
 
@@ -82,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.
@@ -99,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."""
 
@@ -106,21 +142,131 @@ class AbstractMetastore(ABC, Serializable):
         """Cleanup for tests."""
 
     #
-    # Datasets
+    # Namespaces
     #
 
+    @property
+    @abstractmethod
+    def default_namespace_name(self):
+        """Gets default namespace name"""
+
+    @property
+    def system_namespace_name(self):
+        return Namespace.system()
+
+    @abstractmethod
+    def create_namespace(
+        self,
+        name: str,
+        description: str | None = None,
+        uuid: str | None = None,
+        ignore_if_exists: bool = True,
+        validate: bool = True,
+        **kwargs,
+    ) -> Namespace:
+        """Creates new namespace"""
+
+    @abstractmethod
+    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"""
+
+    #
+    # Projects
+    #
+
+    @property
+    @abstractmethod
+    def default_project_name(self):
+        """Gets default project name"""
+
+    @property
+    def listing_project_name(self):
+        return Project.listing()
+
+    @cached_property
+    def default_project(self) -> Project:
+        return self.get_project(
+            self.default_project_name, self.default_namespace_name, create=True
+        )
+
+    @cached_property
+    def listing_project(self) -> Project:
+        return self.get_project(self.listing_project_name, self.system_namespace_name)
+
+    @abstractmethod
+    def create_project(
+        self,
+        namespace_name: str,
+        name: str,
+        description: str | None = None,
+        uuid: str | None = None,
+        ignore_if_exists: bool = True,
+        validate: bool = True,
+        **kwargs,
+    ) -> Project:
+        """Creates new project in specific namespace"""
+
+    @abstractmethod
+    def get_project(
+        self, name: str, namespace_name: str, create: bool = False, conn=None
+    ) -> Project:
+        """
+        Gets a single project inside some namespace by name.
+        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 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)"""
+
+    #
+    # Datasets
+    #
     @abstractmethod
     def create_dataset(
         self,
         name: str,
+        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,
-        labels: Optional[list[str]] = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
     ) -> DatasetRecord:
         """Creates new dataset."""
 
@@ -128,23 +274,23 @@ class AbstractMetastore(ABC, Serializable):
     def create_dataset_version(  # noqa: PLR0913
         self,
         dataset: DatasetRecord,
-        version: int,
+        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."""
 
@@ -158,13 +304,13 @@ class AbstractMetastore(ABC, Serializable):
 
     @abstractmethod
     def update_dataset_version(
-        self, dataset: DatasetRecord, version: int, **kwargs
+        self, dataset: DatasetRecord, version: str, **kwargs
     ) -> DatasetVersion:
         """Updates dataset version fields."""
 
     @abstractmethod
     def remove_dataset_version(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> DatasetRecord:
         """
         Deletes one single dataset version.
@@ -172,15 +318,32 @@ class AbstractMetastore(ABC, Serializable):
         """
 
     @abstractmethod
-    def list_datasets(self) -> Iterator[DatasetListRecord]:
-        """Lists all datasets."""
+    def list_datasets(
+        self, project_id: int | None = None
+    ) -> Iterator[DatasetListRecord]:
+        """Lists all datasets in some project or in all projects."""
 
     @abstractmethod
-    def list_datasets_by_prefix(self, prefix: str) -> Iterator["DatasetListRecord"]:
-        """Lists all datasets which names start with prefix."""
+    def count_datasets(self, project_id: int | None = None) -> int:
+        """Counts datasets in some project or in all projects."""
 
     @abstractmethod
-    def get_dataset(self, name: str) -> DatasetRecord:
+    def list_datasets_by_prefix(
+        self, prefix: str, project_id: int | None = None
+    ) -> Iterator["DatasetListRecord"]:
+        """
+        Lists all datasets which names start with prefix in some project or in all
+        projects.
+        """
+
+    @abstractmethod
+    def get_dataset(
+        self,
+        name: str,  # normal, not full dataset name
+        namespace_name: str | None = None,
+        project_name: str | None = None,
+        conn=None,
+    ) -> DatasetRecord:
         """Gets a single dataset by name."""
 
     @abstractmethod
@@ -188,7 +351,7 @@ class AbstractMetastore(ABC, Serializable):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[int] = None,
+        version: str | None = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -201,10 +364,10 @@ class AbstractMetastore(ABC, Serializable):
     @abstractmethod
     def add_dataset_dependency(
         self,
-        source_dataset_name: str,
-        source_dataset_version: int,
-        dataset_name: str,
-        dataset_version: int,
+        source_dataset: "DatasetRecord",
+        source_dataset_version: str,
+        dep_dataset: "DatasetRecord",
+        dep_dataset_version: str,
     ) -> None:
         """Adds dataset dependency to dataset."""
 
@@ -212,21 +375,27 @@ class AbstractMetastore(ABC, Serializable):
     def update_dataset_dependency_source(
         self,
         source_dataset: DatasetRecord,
-        source_dataset_version: int,
-        new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[int] = None,
+        source_dataset_version: str,
+        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: int
-    ) -> list[Optional[DatasetDependency]]:
+        self, dataset: DatasetRecord, version: str
+    ) -> 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[int] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well.
@@ -234,7 +403,7 @@ class AbstractMetastore(ABC, Serializable):
 
     @abstractmethod
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -254,43 +423,121 @@ class AbstractMetastore(ABC, Serializable):
         name: str,
         query: str,
         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.
         Returns the job id.
         """
 
+    @abstractmethod
+    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: 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
     def set_job_status(
         self,
         job_id: str,
         status: JobStatus,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
-        metrics: Optional[dict[str, Any]] = 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 set_job_and_dataset_status(
+    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,
-        job_status: JobStatus,
-        dataset_status: DatasetStatus,
+        _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:
-        """Set the status of the given job and dataset."""
+        """
+        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_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
-        """Returns dataset names and versions for the job."""
-        raise NotImplementedError
+    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):
@@ -301,14 +548,18 @@ class AbstractDBMetastore(AbstractMetastore):
     and has shared logic for all database systems currently in use.
     """
 
+    NAMESPACE_TABLE = "namespaces"
+    PROJECT_TABLE = "projects"
     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)
 
@@ -319,14 +570,65 @@ class AbstractDBMetastore(AbstractMetastore):
     def cleanup_tables(self, temp_table_names: list[str]) -> None:
         """Cleanup temp tables."""
 
+    @classmethod
+    def _namespaces_columns(cls) -> list["SchemaItem"]:
+        """Namespace table columns."""
+        return [
+            Column("id", Integer, primary_key=True),
+            Column("uuid", Text, nullable=False, default=uuid4()),
+            Column("name", Text, nullable=False),
+            Column("description", Text),
+            Column("created_at", DateTime(timezone=True)),
+        ]
+
+    @cached_property
+    def _namespaces_fields(self) -> list[str]:
+        return [
+            c.name  # type: ignore [attr-defined]
+            for c in self._namespaces_columns()
+            if c.name  # type: ignore [attr-defined]
+        ]
+
+    @classmethod
+    def _projects_columns(cls) -> list["SchemaItem"]:
+        """Project table columns."""
+        return [
+            Column("id", Integer, primary_key=True),
+            Column("uuid", Text, nullable=False, default=uuid4()),
+            Column("name", Text, nullable=False),
+            Column("description", Text),
+            Column("created_at", DateTime(timezone=True)),
+            Column(
+                "namespace_id",
+                Integer,
+                ForeignKey(f"{cls.NAMESPACE_TABLE}.id", ondelete="CASCADE"),
+                nullable=False,
+            ),
+            UniqueConstraint("namespace_id", "name"),
+        ]
+
+    @cached_property
+    def _projects_fields(self) -> list[str]:
+        return [
+            c.name  # type: ignore [attr-defined]
+            for c in self._projects_columns()
+            if c.name  # type: ignore [attr-defined]
+        ]
+
     @classmethod
     def _datasets_columns(cls) -> list["SchemaItem"]:
         """Datasets table columns."""
         return [
             Column("id", Integer, primary_key=True),
+            Column(
+                "project_id",
+                Integer,
+                ForeignKey(f"{cls.PROJECT_TABLE}.id", ondelete="CASCADE"),
+                nullable=False,
+            ),
             Column("name", Text, nullable=False),
             Column("description", Text),
-            Column("labels", JSON, nullable=True),
+            Column("attrs", JSON, nullable=True),
             Column("status", Integer, nullable=False),
             Column("feature_schema", JSON, nullable=True),
             Column("created_at", DateTime(timezone=True)),
@@ -367,7 +669,7 @@ class AbstractDBMetastore(AbstractMetastore):
                 ForeignKey(f"{cls.DATASET_TABLE}.id", ondelete="CASCADE"),
                 nullable=False,
             ),
-            Column("version", Integer, nullable=False),
+            Column("version", Text, nullable=False, default="1.0.0"),
             Column(
                 "status",
                 Integer,
@@ -442,6 +744,16 @@ class AbstractDBMetastore(AbstractMetastore):
     #
     # Query Tables
     #
+    @cached_property
+    def _namespaces(self) -> Table:
+        return Table(
+            self.NAMESPACE_TABLE, self.db.metadata, *self._namespaces_columns()
+        )
+
+    @cached_property
+    def _projects(self) -> Table:
+        return Table(self.PROJECT_TABLE, self.db.metadata, *self._projects_columns())
+
     @cached_property
     def _datasets(self) -> Table:
         return Table(self.DATASET_TABLE, self.db.metadata, *self._datasets_columns())
@@ -465,6 +777,31 @@ class AbstractDBMetastore(AbstractMetastore):
     #
     # Query Starters (These can be overridden by subclasses)
     #
+    @abstractmethod
+    def _namespaces_insert(self) -> "Insert": ...
+
+    def _namespaces_select(self, *columns) -> "Select":
+        if not columns:
+            return self._namespaces.select()
+        return select(*columns)
+
+    def _namespaces_update(self) -> "Update":
+        return self._namespaces.update()
+
+    def _namespaces_delete(self) -> "Delete":
+        return self._namespaces.delete()
+
+    @abstractmethod
+    def _projects_insert(self) -> "Insert": ...
+
+    def _projects_select(self, *columns) -> "Select":
+        if not columns:
+            return self._projects.select()
+        return select(*columns)
+
+    def _projects_delete(self) -> "Delete":
+        return self._projects.delete()
+
     @abstractmethod
     def _datasets_insert(self) -> "Insert": ...
 
@@ -507,6 +844,197 @@ class AbstractDBMetastore(AbstractMetastore):
     def _datasets_dependencies_delete(self) -> "Delete":
         return self._datasets_dependencies.delete()
 
+    #
+    # Namespaces
+    #
+
+    def create_namespace(
+        self,
+        name: str,
+        description: str | None = None,
+        uuid: str | None = None,
+        ignore_if_exists: bool = True,
+        validate: bool = True,
+        **kwargs,
+    ) -> Namespace:
+        if validate:
+            Namespace.validate_name(name)
+        query = self._namespaces_insert().values(
+            name=name,
+            uuid=uuid or str(uuid4()),
+            created_at=datetime.now(timezone.utc),
+            description=description,
+        )
+        if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
+            # SQLite and PostgreSQL both support 'on_conflict_do_nothing',
+            # but generic SQL does not
+            query = query.on_conflict_do_nothing(index_elements=["name"])
+        self.db.execute(query)
+
+        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
+
+        query = self._namespaces_select(
+            *(getattr(n.c, f) for f in self._namespaces_fields),
+        ).where(n.c.name == name)
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            raise NamespaceNotFoundError(f"Namespace {name} not found.")
+        return self.namespace_class.parse(*rows[0])
+
+    def list_namespaces(self, conn=None) -> list[Namespace]:
+        """Gets a list of all namespaces"""
+        n = self._namespaces
+
+        query = self._namespaces_select(
+            *(getattr(n.c, f) for f in self._namespaces_fields),
+        )
+        rows = list(self.db.execute(query, conn=conn))
+
+        return [self.namespace_class.parse(*r) for r in rows]
+
+    #
+    # Projects
+    #
+
+    def create_project(
+        self,
+        namespace_name: str,
+        name: str,
+        description: str | None = None,
+        uuid: str | None = None,
+        ignore_if_exists: bool = True,
+        validate: bool = True,
+        **kwargs,
+    ) -> Project:
+        if validate:
+            Project.validate_name(name)
+        try:
+            namespace = self.get_namespace(namespace_name)
+        except NamespaceNotFoundError:
+            namespace = self.create_namespace(namespace_name, validate=validate)
+
+        query = self._projects_insert().values(
+            namespace_id=namespace.id,
+            uuid=uuid or str(uuid4()),
+            name=name,
+            created_at=datetime.now(timezone.utc),
+            description=description,
+        )
+        if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
+            # SQLite and PostgreSQL both support 'on_conflict_do_nothing',
+            # but generic SQL does not
+            query = query.on_conflict_do_nothing(
+                index_elements=["namespace_id", "name"]
+            )
+        self.db.execute(query)
+
+        return self.get_project(name, namespace.name)
+
+    def _projects_base_query(self) -> "Select":
+        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),
+        )
+        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
+    ) -> Project:
+        """Gets a single project inside some namespace by name"""
+        n = self._namespaces
+        p = self._projects
+        validate = True
+
+        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_base_query().where(
+            p.c.name == name, n.c.name == namespace_name
+        )
+
+        rows = list(self.db.execute(query, conn=conn))
+        if not rows:
+            if create:
+                return self.create_project(namespace_name, name, validate=validate)
+            raise ProjectNotFoundError(
+                f"Project {name} in namespace {namespace_name} not found."
+            )
+        return self.project_class.parse(*rows[0])
+
+    def get_project_by_id(self, project_id: int, conn=None) -> Project:
+        """Gets a single project by id"""
+        p = self._projects
+
+        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 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
+        """
+        p = self._projects
+
+        query = self._projects_base_query()
+
+        if namespace_id:
+            query = query.where(p.c.namespace_id == namespace_id)
+
+        rows = list(self.db.execute(query, conn=conn))
+
+        return [self.project_class.parse(*r) for r in rows]
+
     #
     # Datasets
     #
@@ -514,20 +1042,26 @@ class AbstractDBMetastore(AbstractMetastore):
     def create_dataset(
         self,
         name: str,
+        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,
-        labels: Optional[list[str]] = None,
+        description: str | None = None,
+        attrs: list[str] | None = None,
         **kwargs,  # TODO registered = True / False
     ) -> DatasetRecord:
         """Creates new dataset."""
-        # TODO abstract this method and add registered = True based on kwargs
+        if not project_id:
+            project = self.default_project
+        else:
+            project = self.get_project_by_id(project_id)
+
         query = self._datasets_insert().values(
             name=name,
+            project_id=project.id,
             status=status,
             feature_schema=json.dumps(feature_schema or {}),
             created_at=datetime.now(timezone.utc),
@@ -538,36 +1072,38 @@ class AbstractDBMetastore(AbstractMetastore):
             query_script=query_script,
             schema=json.dumps(schema or {}),
             description=description,
-            labels=json.dumps(labels or []),
+            attrs=json.dumps(attrs or []),
         )
         if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
             # SQLite and PostgreSQL both support 'on_conflict_do_nothing',
             # but generic SQL does not
-            query = query.on_conflict_do_nothing(index_elements=["name"])
+            query = query.on_conflict_do_nothing(index_elements=["project_id", "name"])
         self.db.execute(query)
 
-        return self.get_dataset(name)
+        return self.get_dataset(
+            name, namespace_name=project.namespace.name, project_name=project.name
+        )
 
     def create_dataset_version(  # noqa: PLR0913
         self,
         dataset: DatasetRecord,
-        version: int,
+        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."""
@@ -603,7 +1139,12 @@ class AbstractDBMetastore(AbstractMetastore):
             )
         self.db.execute(query, conn=conn)
 
-        return self.get_dataset(dataset.name, conn=conn)
+        return self.get_dataset(
+            dataset.name,
+            namespace_name=dataset.project.namespace.name,
+            project_name=dataset.project.name,
+            conn=conn,
+        )
 
     def remove_dataset(self, dataset: DatasetRecord) -> None:
         """Removes dataset."""
@@ -617,26 +1158,47 @@ class AbstractDBMetastore(AbstractMetastore):
         self, dataset: DatasetRecord, conn=None, **kwargs
     ) -> DatasetRecord:
         """Updates dataset fields."""
-        values = {}
-        dataset_values = {}
+        values: dict[str, Any] = {}
+        dataset_values: dict[str, Any] = {}
         for field, value in kwargs.items():
-            if field in self._dataset_fields[1:]:
-                if field in ["labels", "schema"]:
-                    values[field] = json.dumps(value) if value else None
+            if field in ("id", "created_at") or field not in self._dataset_fields:
+                continue  # these fields are read-only or not applicable
+
+            if value is None and field in ("name", "status", "sources", "query_script"):
+                raise ValueError(f"Field {field} cannot be None")
+            if field == "name" and not value:
+                raise ValueError("name cannot be empty")
+
+            if field == "attrs":
+                if value is None:
+                    values[field] = None
                 else:
-                    values[field] = value
-                if field == "schema":
-                    dataset_values[field] = DatasetRecord.parse_schema(value)
+                    values[field] = json.dumps(value)
+                dataset_values[field] = value
+            elif field == "schema":
+                if value is None:
+                    values[field] = None
+                    dataset_values[field] = None
                 else:
-                    dataset_values[field] = value
+                    values[field] = json.dumps(value)
+                    dataset_values[field] = parse_schema(value)
+            elif field == "project_id":
+                if not value:
+                    raise ValueError("Cannot set empty project_id for dataset")
+                dataset_values["project"] = self.get_project_by_id(value)
+                values[field] = value
+            else:
+                values[field] = value
+                dataset_values[field] = value
 
         if not values:
-            # Nothing to update
-            return dataset
+            return dataset  # nothing to update
 
         d = self._datasets
         self.db.execute(
-            self._datasets_update().where(d.c.name == dataset.name).values(values),
+            self._datasets_update()
+            .where(d.c.name == dataset.name, d.c.project_id == dataset.project.id)
+            .values(values),
             conn=conn,
         )  # type: ignore [attr-defined]
 
@@ -645,46 +1207,79 @@ class AbstractDBMetastore(AbstractMetastore):
         return result_ds
 
     def update_dataset_version(
-        self, dataset: DatasetRecord, version: int, conn=None, **kwargs
+        self, dataset: DatasetRecord, version: str, conn=None, **kwargs
     ) -> DatasetVersion:
         """Updates dataset fields."""
-        dataset_version = dataset.get_version(version)
-
-        values = {}
+        values: dict[str, Any] = {}
+        version_values: dict[str, Any] = {}
         for field, value in kwargs.items():
-            if field in self._dataset_version_fields[1:]:
-                if field == "schema":
-                    dataset_version.update(**{field: DatasetRecord.parse_schema(value)})
-                    values[field] = json.dumps(value) if value else None
-                elif field == "feature_schema":
-                    values[field] = json.dumps(value) if value else None
-                elif field == "preview" and isinstance(value, list):
-                    values[field] = json.dumps(value, cls=JSONSerialize)
+            if (
+                field in ("id", "created_at")
+                or field not in self._dataset_version_fields
+            ):
+                continue  # these fields are read-only or not applicable
+
+            if value is None and field in (
+                "status",
+                "sources",
+                "query_script",
+                "error_message",
+                "error_stack",
+                "script_output",
+                "uuid",
+            ):
+                raise ValueError(f"Field {field} cannot be None")
+
+            if field == "schema":
+                values[field] = json.dumps(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
+                else:
+                    values[field] = json.dumps(value)
+                version_values[field] = value
+            elif field == "preview":
+                if value is None:
+                    values[field] = None
+                elif not isinstance(value, list):
+                    raise ValueError(
+                        f"Field '{field}' must be a list, got {type(value).__name__}"
+                    )
                 else:
-                    values[field] = value
-                    dataset_version.update(**{field: value})
+                    values[field] = json.dumps(value, serialize_bytes=True)
+                version_values["_preview_data"] = value
+            else:
+                values[field] = value
+                version_values[field] = value
 
         if not values:
-            # Nothing to update
-            return dataset_version
+            return dataset.get_version(version)
 
         dv = self._datasets_versions
         self.db.execute(
             self._datasets_versions_update()
-            .where(dv.c.id == dataset_version.id)
+            .where(dv.c.dataset_id == dataset.id, dv.c.version == version)
             .values(values),
             conn=conn,
         )  # type: ignore [attr-defined]
 
-        return dataset_version
+        for v in dataset.versions:
+            if v.version == version:
+                v.update(**version_values)
+                return v
 
-    def _parse_dataset(self, rows) -> Optional[DatasetRecord]:
+        raise DatasetVersionNotFoundError(
+            f"Dataset {dataset.name} does not have version {version}"
+        )
+
+    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
@@ -692,69 +1287,124 @@ class AbstractDBMetastore(AbstractMetastore):
 
     def _parse_dataset_list(self, rows) -> Iterator["DatasetListRecord"]:
         # grouping rows by dataset id
-        for _, g in groupby(rows, lambda r: r[0]):
+        for _, g in groupby(rows, lambda r: r[11]):
             dataset = self._parse_list_dataset(list(g))
             if dataset:
                 yield dataset
 
     def _get_dataset_query(
         self,
+        namespace_fields: list[str],
+        project_fields: list[str],
         dataset_fields: list[str],
         dataset_version_fields: list[str],
         isouter: bool = True,
-    ):
+    ) -> "Select":
         if not (
             self.db.has_table(self._datasets.name)
             and self.db.has_table(self._datasets_versions.name)
         ):
             raise TableMissingError
 
+        n = self._namespaces
+        p = self._projects
         d = self._datasets
         dv = self._datasets_versions
 
         query = self._datasets_select(
+            *(getattr(n.c, f) for f in namespace_fields),
+            *(getattr(p.c, f) for f in project_fields),
             *(getattr(d.c, f) for f in dataset_fields),
             *(getattr(dv.c, f) for f in dataset_version_fields),
         )
-        j = d.join(dv, d.c.id == dv.c.dataset_id, isouter=isouter)
+        j = (
+            n.join(p, n.c.id == p.c.namespace_id)
+            .join(d, p.c.id == d.c.project_id)
+            .join(dv, d.c.id == dv.c.dataset_id, isouter=isouter)
+        )
         return query.select_from(j)
 
-    def _base_dataset_query(self):
+    def _base_dataset_query(self) -> "Select":
         return self._get_dataset_query(
-            self._dataset_fields, self._dataset_version_fields
+            self._namespaces_fields,
+            self._projects_fields,
+            self._dataset_fields,
+            self._dataset_version_fields,
         )
 
-    def _base_list_datasets_query(self):
+    def _base_list_datasets_query(self) -> "Select":
         return self._get_dataset_query(
-            self._dataset_list_fields, self._dataset_list_version_fields, isouter=False
+            self._namespaces_fields,
+            self._projects_fields,
+            self._dataset_list_fields,
+            self._dataset_list_version_fields,
+            isouter=False,
         )
 
-    def list_datasets(self) -> Iterator["DatasetListRecord"]:
-        """Lists all datasets."""
+    def list_datasets(
+        self, project_id: int | None = None
+    ) -> Iterator["DatasetListRecord"]:
+        d = self._datasets
         query = self._base_list_datasets_query().order_by(
             self._datasets.c.name, self._datasets_versions.c.version
         )
+        if project_id:
+            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, conn=None
+        self, prefix: str, project_id: int | None = None, conn=None
     ) -> Iterator["DatasetListRecord"]:
+        d = self._datasets
         query = self._base_list_datasets_query()
+        if project_id:
+            query = query.where(d.c.project_id == project_id)
         query = query.where(self._datasets.c.name.startswith(prefix))
         yield from self._parse_dataset_list(self.db.execute(query))
 
-    def get_dataset(self, name: str, conn=None) -> DatasetRecord:
-        """Gets a single dataset by name"""
+    def get_dataset(
+        self,
+        name: str,  # normal, not full dataset name
+        namespace_name: str | None = None,
+        project_name: str | None = None,
+        conn=None,
+    ) -> DatasetRecord:
+        """
+        Gets a single dataset in project by dataset name.
+        """
+        namespace_name = namespace_name or self.default_namespace_name
+        project_name = project_name or self.default_project_name
+
         d = self._datasets
+        n = self._namespaces
+        p = self._projects
         query = self._base_dataset_query()
-        query = query.where(d.c.name == name)  # type: ignore [attr-defined]
+        query = query.where(
+            d.c.name == name,
+            n.c.name == namespace_name,
+            p.c.name == project_name,
+        )  # type: ignore [attr-defined]
         ds = self._parse_dataset(self.db.execute(query, conn=conn))
         if not ds:
-            raise DatasetNotFoundError(f"Dataset {name} not found.")
+            raise DatasetNotFoundError(
+                f"Dataset {name} not found in namespace {namespace_name}"
+                f" and project {project_name}"
+            )
+
         return ds
 
     def remove_dataset_version(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> DatasetRecord:
         """
         Deletes one single dataset version.
@@ -787,7 +1437,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[int] = None,
+        version: str | None = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -808,7 +1458,7 @@ class AbstractDBMetastore(AbstractMetastore):
             update_data["error_message"] = error_message
             update_data["error_stack"] = error_stack
 
-        self.update_dataset(dataset, conn=conn, **update_data)
+        dataset = self.update_dataset(dataset, conn=conn, **update_data)
 
         if version:
             self.update_dataset_version(dataset, version, conn=conn, **update_data)
@@ -820,32 +1470,29 @@ class AbstractDBMetastore(AbstractMetastore):
     #
     def add_dataset_dependency(
         self,
-        source_dataset_name: str,
-        source_dataset_version: int,
-        dataset_name: str,
-        dataset_version: int,
+        source_dataset: "DatasetRecord",
+        source_dataset_version: str,
+        dep_dataset: "DatasetRecord",
+        dep_dataset_version: str,
     ) -> None:
         """Adds dataset dependency to dataset."""
-        source_dataset = self.get_dataset(source_dataset_name)
-        dataset = self.get_dataset(dataset_name)
-
         self.db.execute(
             self._datasets_dependencies_insert().values(
                 source_dataset_id=source_dataset.id,
                 source_dataset_version_id=(
                     source_dataset.get_version(source_dataset_version).id
                 ),
-                dataset_id=dataset.id,
-                dataset_version_id=dataset.get_version(dataset_version).id,
+                dataset_id=dep_dataset.id,
+                dataset_version_id=dep_dataset.get_version(dep_dataset_version).id,
             )
         )
 
     def update_dataset_dependency_source(
         self,
         source_dataset: DatasetRecord,
-        source_dataset_version: int,
-        new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[int] = None,
+        source_dataset_version: str,
+        new_source_dataset: DatasetRecord | None = None,
+        new_source_dataset_version: str | None = None,
     ) -> None:
         dd = self._datasets_dependencies
 
@@ -875,9 +1522,23 @@ 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: int
-    ) -> list[Optional[DatasetDependency]]:
+        self, dataset: DatasetRecord, version: str
+    ) -> list[DatasetDependency | None]:
+        n = self._namespaces
+        p = self._projects
         d = self._datasets
         dd = self._datasets_dependencies
         dv = self._datasets_versions
@@ -889,23 +1550,90 @@ class AbstractDBMetastore(AbstractMetastore):
         query = (
             self._datasets_dependencies_select(*select_cols)
             .select_from(
-                dd.join(d, dd.c.dataset_id == d.c.id, isouter=True).join(
-                    dv, dd.c.dataset_version_id == dv.c.id, isouter=True
-                )
+                dd.join(d, dd.c.dataset_id == d.c.id, isouter=True)
+                .join(dv, dd.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)
             )
             .where(
                 (dd.c.source_dataset_id == dataset.id)
                 & (dd.c.source_dataset_version_id == dataset_version.id)
             )
         )
-        if version:
-            dataset_version = dataset.get_version(version)
-            query = query.where(dd.c.source_dataset_version_id == dataset_version.id)
 
         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[int] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well
@@ -924,7 +1652,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self.db.execute(q)
 
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: str | None = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -975,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":
@@ -1013,15 +1743,29 @@ 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,
         query: str,
         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.
@@ -1032,7 +1776,7 @@ class AbstractDBMetastore(AbstractMetastore):
             self._jobs_insert().values(
                 id=job_id,
                 name=name,
-                status=JobStatus.CREATED,
+                status=status,
                 created_at=datetime.now(timezone.utc),
                 query=query,
                 query_type=query_type.value,
@@ -1042,30 +1786,68 @@ 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) -> 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))
+        if not results:
+            return None
+        return self._parse_job(results[0])
+
+    def update_job(
+        self,
+        job_id: str,
+        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:
+            values["status"] = status
+        if error_message is not None:
+            values["error_message"] = error_message
+        if error_stack is not None:
+            values["error_stack"] = error_stack
+        if finished_at is not None:
+            values["finished_at"] = finished_at
+        if metrics:
+            values["metrics"] = json.dumps(metrics)
+
+        if values:
+            j = self._jobs
+            self.db.execute(
+                self._jobs_update().where(j.c.id == job_id).values(**values),
+                conn=conn,
+            )  # type: ignore [attr-defined]
+
+        return self.get_job(job_id, conn=conn)
+
     def set_job_status(
         self,
         job_id: str,
         status: JobStatus,
-        error_message: Optional[str] = None,
-        error_stack: Optional[str] = None,
-        metrics: Optional[dict[str, Any]] = 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.value}
-        if status.value in JobStatus.finished():
+        values: dict = {"status": status}
+        if status in JobStatus.finished():
             values["finished_at"] = datetime.now(timezone.utc)
         if error_message:
             values["error_message"] = error_message
         if error_stack:
             values["error_stack"] = error_stack
-        if metrics:
-            values["metrics"] = json.dumps(metrics)
         self.db.execute(
             self._jobs_update(self._jobs.c.id == job_id).values(**values),
             conn=conn,
@@ -1074,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(
@@ -1087,36 +1869,320 @@ class AbstractDBMetastore(AbstractMetastore):
             return None
         return results[0][0]
 
-    def set_job_and_dataset_status(
+    #
+    # 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,
-        job_status: JobStatus,
-        dataset_status: DatasetStatus,
+        is_creator: bool = False,
+        conn=None,
     ) -> None:
-        """Set the status of the given job and dataset."""
-        with self.db.transaction() as conn:
-            self.set_job_status(job_id, status=job_status, conn=conn)
-            dv = self._datasets_versions
-            query = (
-                self._datasets_versions_update()
-                .where(
-                    (dv.c.job_id == job_id) & (dv.c.status != DatasetStatus.COMPLETE)
+        # 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"]
                 )
-                .values(status=dataset_status)
+            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"),
             )
-            self.db.execute(query, conn=conn)  # type: ignore[attr-defined]
+            .where(self._jobs.c.id == job_id)
+            .cte(name="ancestors", recursive=True)
+        )
 
-    def get_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
-        """Returns dataset names and versions for the job."""
-        dv = self._datasets_versions
-        ds = self._datasets
+        # 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
+        )
 
-        join_condition = dv.c.dataset_id == ds.c.id
+        results = list(self.db.execute(query, conn=conn))
 
-        query = (
-            self._datasets_versions_select(ds.c.name, dv.c.version)
-            .select_from(dv.join(ds, join_condition))
-            .where(dv.c.job_id == job_id)
+        # 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)
         )
 
-        return list(self.db.execute(query))
+    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])