datachain 0.18.3__py3-none-any.whl → 0.18.5__py3-none-any.whl

datachain/catalog/catalog.py

@@ -66,6 +66,7 @@ if TYPE_CHECKING:
     )
     from datachain.dataset import DatasetListVersion
     from datachain.job import Job
+    from datachain.lib.listing_info import ListingInfo
     from datachain.listing import Listing
 
 logger = logging.getLogger("datachain")
@@ -910,11 +911,7 @@ class Catalog:
             values["num_objects"] = None
             values["size"] = None
             values["preview"] = None
-            self.metastore.update_dataset_version(
-                dataset,
-                version,
-                **values,
-            )
+            self.metastore.update_dataset_version(dataset, version, **values)
             return
 
         if not dataset_version.num_objects:
@@ -934,11 +931,7 @@ class Catalog:
         if not values:
             return
 
-        self.metastore.update_dataset_version(
-            dataset,
-            version,
-            **values,
-        )
+        self.metastore.update_dataset_version(dataset, version, **values)
 
     def update_dataset(
         self, dataset: DatasetRecord, conn=None, **kwargs
@@ -1116,13 +1109,16 @@ class Catalog:
         return direct_dependencies
 
     def ls_datasets(
-        self, include_listing: bool = False, studio: bool = False
+        self,
+        prefix: Optional[str] = None,
+        include_listing: bool = False,
+        studio: bool = False,
     ) -> Iterator[DatasetListRecord]:
         from datachain.remote.studio import StudioClient
 
         if studio:
             client = StudioClient()
-            response = client.ls_datasets()
+            response = client.ls_datasets(prefix=prefix)
             if not response.ok:
                 raise DataChainError(response.message)
             if not response.data:
@@ -1133,6 +1129,8 @@ class Catalog:
                 for d in response.data
                 if not d.get("name", "").startswith(QUERY_DATASET_PREFIX)
             )
+        elif prefix:
+            datasets = self.metastore.list_datasets_by_prefix(prefix)
         else:
             datasets = self.metastore.list_datasets()
 
@@ -1142,39 +1140,55 @@ class Catalog:
 
     def list_datasets_versions(
         self,
+        prefix: Optional[str] = None,
         include_listing: bool = False,
+        with_job: bool = True,
         studio: bool = False,
     ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", Optional["Job"]]]:
         """Iterate over all dataset versions with related jobs."""
         datasets = list(
-            self.ls_datasets(include_listing=include_listing, studio=studio)
+            self.ls_datasets(
+                prefix=prefix, include_listing=include_listing, studio=studio
+            )
         )
 
         # preselect dataset versions jobs from db to avoid multiple queries
-        jobs_ids: set[str] = {
-            v.job_id for ds in datasets for v in ds.versions if v.job_id
-        }
         jobs: dict[str, Job] = {}
-        if jobs_ids:
-            jobs = {j.id: j for j in self.metastore.list_jobs_by_ids(list(jobs_ids))}
+        if with_job:
+            jobs_ids: set[str] = {
+                v.job_id for ds in datasets for v in ds.versions if v.job_id
+            }
+            if jobs_ids:
+                jobs = {
+                    j.id: j for j in self.metastore.list_jobs_by_ids(list(jobs_ids))
+                }
 
         for d in datasets:
             yield from (
-                (d, v, jobs.get(str(v.job_id)) if v.job_id else None)
+                (d, v, jobs.get(str(v.job_id)) if with_job and v.job_id else None)
                 for v in d.versions
             )
 
-    def listings(self):
+    def listings(self, prefix: Optional[str] = None) -> list["ListingInfo"]:
         """
         Returns list of ListingInfo objects which are representing specific
         storage listing datasets
         """
-        from datachain.lib.listing import is_listing_dataset
+        from datachain.lib.listing import LISTING_PREFIX, is_listing_dataset
         from datachain.lib.listing_info import ListingInfo
 
+        if prefix and not prefix.startswith(LISTING_PREFIX):
+            prefix = LISTING_PREFIX + prefix
+
+        listing_datasets_versions = self.list_datasets_versions(
+            prefix=prefix,
+            include_listing=True,
+            with_job=False,
+        )
+
         return [
             ListingInfo.from_models(d, v, j)
-            for d, v, j in self.list_datasets_versions(include_listing=True)
+            for d, v, j in listing_datasets_versions
             if is_listing_dataset(d.name)
         ]
 

datachain/client/azure.py

@@ -65,7 +65,7 @@ class AzureClient(Client):
                         if entries:
                             await result_queue.put(entries)
                             pbar.update(len(entries))
-                    if not found:
+                    if not found and prefix:
                         raise FileNotFoundError(
                             f"Unable to resolve remote path: {prefix}"
                         )

datachain/client/gcs.py

@@ -74,7 +74,7 @@ class GCSClient(Client):
             try:
                 await self._get_pages(prefix, page_queue)
                 found = await consumer
-                if not found:
+                if not found and prefix:
                     raise FileNotFoundError(f"Unable to resolve remote path: {prefix}")
             finally:
                 consumer.cancel()  # In case _get_pages() raised

datachain/client/s3.py

@@ -80,7 +80,7 @@ class ClientS3(Client):
             finally:
                 await page_queue.put(None)
 
-        async def process_pages(page_queue, result_queue):
+        async def process_pages(page_queue, result_queue, prefix):
             found = False
             with tqdm(desc=f"Listing {self.uri}", unit=" objects", leave=False) as pbar:
                 while (res := await page_queue.get()) is not None:
@@ -94,7 +94,7 @@ class ClientS3(Client):
                     if entries:
                         await result_queue.put(entries)
                         pbar.update(len(entries))
-            if not found:
+            if not found and prefix:
                 raise FileNotFoundError(f"Unable to resolve remote path: {prefix}")
 
         try:
@@ -118,7 +118,9 @@ class ClientS3(Client):
                 Delimiter="",
             )
             page_queue: asyncio.Queue[list] = asyncio.Queue(2)
-            consumer = asyncio.create_task(process_pages(page_queue, result_queue))
+            consumer = asyncio.create_task(
+                process_pages(page_queue, result_queue, prefix)
+            )
             try:
                 await get_pages(it, page_queue)
                 await consumer

datachain/data_storage/metastore.py

@@ -36,6 +36,7 @@ from datachain.dataset import (
 )
 from datachain.error import (
     DatasetNotFoundError,
+    DatasetVersionNotFoundError,
     TableMissingError,
 )
 from datachain.job import Job
@@ -273,7 +274,6 @@ class AbstractMetastore(ABC, Serializable):
         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,
@@ -620,22 +620,36 @@ 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 ["attrs", "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] = DatasetRecord.parse_schema(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(
@@ -651,36 +665,70 @@ class AbstractDBMetastore(AbstractMetastore):
         self, dataset: DatasetRecord, version: str, conn=None, **kwargs
     ) -> DatasetVersion:
         """Updates dataset fields."""
-        dataset_version = dataset.get_version(version)
-
-        values = {}
-        version_values: dict = {}
+        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":
-                    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
+            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] = (
+                    DatasetRecord.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
-                    version_values[field] = value
+                    values[field] = json.dumps(value, cls=JSONSerialize)
+                version_values["_preview_data"] = value
+            else:
+                values[field] = value
+                version_values[field] = value
 
-        if values:
-            dv = self._datasets_versions
-            self.db.execute(
-                self._datasets_versions_update()
-                .where(dv.c.dataset_id == dataset.id, dv.c.version == version)
-                .values(values),
-                conn=conn,
-            )  # type: ignore [attr-defined]
-            dataset_version.update(**version_values)
+        if not values:
+            return dataset.get_version(version)
+
+        dv = self._datasets_versions
+        self.db.execute(
+            self._datasets_versions_update()
+            .where(dv.c.dataset_id == dataset.id, dv.c.version == version)
+            .values(values),
+            conn=conn,
+        )  # type: ignore [attr-defined]
+
+        for v in dataset.versions:
+            if v.version == version:
+                v.update(**version_values)
+                return v
 
-        return dataset_version
+        raise DatasetVersionNotFoundError(
+            f"Dataset {dataset.name} does not have version {version}"
+        )
 
     def _parse_dataset(self, rows) -> Optional[DatasetRecord]:
         versions = [self.dataset_class.parse(*r) for r in rows]
@@ -812,7 +860,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)
@@ -1064,7 +1112,6 @@ class AbstractDBMetastore(AbstractMetastore):
         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,
@@ -1075,8 +1122,6 @@ class AbstractDBMetastore(AbstractMetastore):
         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:

datachain/dataset.py

@@ -93,7 +93,7 @@ class DatasetDependency:
         if self.type == DatasetDependencyType.DATASET:
             return self.name
 
-        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"), {})
+        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"))
         assert list_dataset_name
         return list_dataset_name