datachain 0.16.4__py3-none-any.whl → 0.17.0__py3-none-any.whl

datachain/data_storage/metastore.py

@@ -128,7 +128,7 @@ 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,
@@ -158,13 +158,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.
@@ -188,7 +188,7 @@ class AbstractMetastore(ABC, Serializable):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[int] = None,
+        version: Optional[str] = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -202,9 +202,9 @@ class AbstractMetastore(ABC, Serializable):
     def add_dataset_dependency(
         self,
         source_dataset_name: str,
-        source_dataset_version: int,
+        source_dataset_version: str,
         dataset_name: str,
-        dataset_version: int,
+        dataset_version: str,
     ) -> None:
         """Adds dataset dependency to dataset."""
 
@@ -212,21 +212,21 @@ class AbstractMetastore(ABC, Serializable):
     def update_dataset_dependency_source(
         self,
         source_dataset: DatasetRecord,
-        source_dataset_version: int,
+        source_dataset_version: str,
         new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[int] = None,
+        new_source_dataset_version: Optional[str] = None,
     ) -> None:
         """Updates dataset dependency source."""
 
     @abstractmethod
     def get_direct_dataset_dependencies(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[Optional[DatasetDependency]]:
         """Gets direct dataset dependencies."""
 
     @abstractmethod
     def remove_dataset_dependencies(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: Optional[str] = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well.
@@ -234,7 +234,7 @@ class AbstractMetastore(ABC, Serializable):
 
     @abstractmethod
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: Optional[str] = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -254,6 +254,7 @@ 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,
@@ -264,33 +265,35 @@ class AbstractMetastore(ABC, Serializable):
         """
 
     @abstractmethod
-    def set_job_status(
+    def get_job(self, job_id: str) -> Optional[Job]:
+        """Returns the job with the given ID."""
+
+    @abstractmethod
+    def update_job(
         self,
         job_id: str,
-        status: JobStatus,
+        status: Optional[JobStatus] = None,
+        exit_code: Optional[int] = None,
         error_message: Optional[str] = None,
         error_stack: Optional[str] = None,
+        finished_at: Optional[datetime] = None,
         metrics: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Set the status of the given job."""
+    ) -> Optional["Job"]:
+        """Updates job fields."""
 
     @abstractmethod
-    def get_job_status(self, job_id: str) -> Optional[JobStatus]:
-        """Returns the status of the given job."""
-
-    @abstractmethod
-    def set_job_and_dataset_status(
+    def set_job_status(
         self,
         job_id: str,
-        job_status: JobStatus,
-        dataset_status: DatasetStatus,
+        status: JobStatus,
+        error_message: Optional[str] = None,
+        error_stack: Optional[str] = None,
     ) -> None:
-        """Set the status of the given job and dataset."""
+        """Set the status of the given 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_job_status(self, job_id: str) -> Optional[JobStatus]:
+        """Returns the status of the given job."""
 
 
 class AbstractDBMetastore(AbstractMetastore):
@@ -367,7 +370,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,
@@ -551,7 +554,7 @@ class AbstractDBMetastore(AbstractMetastore):
     def create_dataset_version(  # noqa: PLR0913
         self,
         dataset: DatasetRecord,
-        version: int,
+        version: str,
         status: int,
         sources: str = "",
         feature_schema: Optional[dict] = None,
@@ -645,36 +648,37 @@ 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 = {}
+        version_values: dict = {}
         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
+                    version_values[field] = DatasetRecord.parse_schema(value)
                 elif field == "feature_schema":
                     values[field] = json.dumps(value) if value else None
+                    version_values[field] = value
                 elif field == "preview" and isinstance(value, list):
                     values[field] = json.dumps(value, cls=JSONSerialize)
+                    version_values[field] = value
                 else:
                     values[field] = value
-                    dataset_version.update(**{field: value})
-
-        if not values:
-            # Nothing to update
-            return dataset_version
+                    version_values[field] = value
 
-        dv = self._datasets_versions
-        self.db.execute(
-            self._datasets_versions_update()
-            .where(dv.c.id == dataset_version.id)
-            .values(values),
-            conn=conn,
-        )  # type: ignore [attr-defined]
+        if values:
+            dv = self._datasets_versions
+            self.db.execute(
+                self._datasets_versions_update()
+                .where(dv.c.dataset_id == dataset.id and dv.c.version == version)
+                .values(values),
+                conn=conn,
+            )  # type: ignore [attr-defined]
+            dataset_version.update(**version_values)
 
         return dataset_version
 
@@ -702,7 +706,7 @@ class AbstractDBMetastore(AbstractMetastore):
         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)
@@ -719,12 +723,12 @@ class AbstractDBMetastore(AbstractMetastore):
         j = d.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
         )
 
-    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
         )
@@ -754,7 +758,7 @@ class AbstractDBMetastore(AbstractMetastore):
         return ds
 
     def remove_dataset_version(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> DatasetRecord:
         """
         Deletes one single dataset version.
@@ -787,7 +791,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         dataset: DatasetRecord,
         status: int,
-        version: Optional[int] = None,
+        version: Optional[str] = None,
         error_message="",
         error_stack="",
         script_output="",
@@ -821,9 +825,9 @@ class AbstractDBMetastore(AbstractMetastore):
     def add_dataset_dependency(
         self,
         source_dataset_name: str,
-        source_dataset_version: int,
+        source_dataset_version: str,
         dataset_name: str,
-        dataset_version: int,
+        dataset_version: str,
     ) -> None:
         """Adds dataset dependency to dataset."""
         source_dataset = self.get_dataset(source_dataset_name)
@@ -843,9 +847,9 @@ class AbstractDBMetastore(AbstractMetastore):
     def update_dataset_dependency_source(
         self,
         source_dataset: DatasetRecord,
-        source_dataset_version: int,
+        source_dataset_version: str,
         new_source_dataset: Optional[DatasetRecord] = None,
-        new_source_dataset_version: Optional[int] = None,
+        new_source_dataset_version: Optional[str] = None,
     ) -> None:
         dd = self._datasets_dependencies
 
@@ -876,7 +880,7 @@ class AbstractDBMetastore(AbstractMetastore):
         """
 
     def get_direct_dataset_dependencies(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[Optional[DatasetDependency]]:
         d = self._datasets
         dd = self._datasets_dependencies
@@ -905,7 +909,7 @@ class AbstractDBMetastore(AbstractMetastore):
         return [self.dependency_class.parse(*r) for r in self.db.execute(query)]
 
     def remove_dataset_dependencies(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: Optional[str] = None
     ) -> None:
         """
         When we remove dataset, we need to clean up it's dependencies as well
@@ -924,7 +928,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self.db.execute(q)
 
     def remove_dataset_dependants(
-        self, dataset: DatasetRecord, version: Optional[int] = None
+        self, dataset: DatasetRecord, version: Optional[str] = None
     ) -> None:
         """
         When we remove dataset, we need to clear its references in other dataset
@@ -1018,6 +1022,7 @@ class AbstractDBMetastore(AbstractMetastore):
         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,
@@ -1032,7 +1037,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,
@@ -1047,25 +1052,65 @@ class AbstractDBMetastore(AbstractMetastore):
         )
         return job_id
 
+    def get_job(self, job_id: str, conn=None) -> Optional[Job]:
+        """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: Optional[JobStatus] = None,
+        exit_code: Optional[int] = 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"]:
+        """Updates job fields."""
+        values: dict = {}
+        if status is not None:
+            values["status"] = status
+        if exit_code is not None:
+            values["exit_code"] = exit_code
+        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,
     ) -> 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,
@@ -1086,37 +1131,3 @@ class AbstractDBMetastore(AbstractMetastore):
         if not results:
             return None
         return results[0][0]
-
-    def set_job_and_dataset_status(
-        self,
-        job_id: str,
-        job_status: JobStatus,
-        dataset_status: DatasetStatus,
-    ) -> 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)
-                )
-                .values(status=dataset_status)
-            )
-            self.db.execute(query, conn=conn)  # type: ignore[attr-defined]
-
-    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
-
-        join_condition = dv.c.dataset_id == ds.c.id
-
-        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)
-        )
-
-        return list(self.db.execute(query))

datachain/data_storage/sqlite.py

@@ -25,6 +25,7 @@ from sqlalchemy.sql.selectable import Select
 from tqdm.auto import tqdm
 
 import datachain.sql.sqlite
+from datachain import semver
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
 from datachain.data_storage.schema import DefaultSchema
@@ -486,7 +487,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         return table
 
     def get_dataset_sources(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[StorageURI]:
         dr = self.dataset_rows(dataset, version)
         query = dr.select(dr.c("source", column="file")).distinct()
@@ -502,8 +503,8 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         src: DatasetRecord,
         dst: DatasetRecord,
-        src_version: int,
-        dst_version: int,
+        src_version: str,
+        dst_version: str,
     ) -> None:
         dst_empty = False
 
@@ -534,7 +535,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             dst_previous_versions = [
                 v.version
                 for v in dst.versions  # type: ignore [union-attr]
-                if v.version < dst_version
+                if semver.compare(v.version, dst_version) == -1
             ]
             if dst_previous_versions:
                 dst_version_latest = max(dst_previous_versions)
@@ -570,7 +571,7 @@ class SQLiteWarehouse(AbstractWarehouse):
                 conn=conn,
             )
 
-    def insert_dataset_rows(self, df, dataset: DatasetRecord, version: int) -> int:
+    def insert_dataset_rows(self, df, dataset: DatasetRecord, version: str) -> int:
         dr = self.dataset_rows(dataset, version)
         return self.db.insert_dataframe(dr.table.name, df)
 
@@ -595,7 +596,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         return col_type.python_type
 
     def dataset_table_export_file_names(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[str]:
         raise NotImplementedError("Exporting dataset table not implemented for SQLite")
 
@@ -603,7 +604,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         bucket_uri: str,
         dataset: DatasetRecord,
-        version: int,
+        version: str,
         client_config=None,
     ) -> list[str]:
         raise NotImplementedError("Exporting dataset table not implemented for SQLite")