get-dataset metastore: exclude preview and versions by default (#1661)

* get-dataset metastore: exclude preview

* add filter to exclude versions

* more cleanup after review

* fix listing dataset fetching

* fix tests: update datasets and tests

* fix listing requires versions

* adjust test fixtures

* fix more places where we need versions after all

* fix tests: explicitily require version where it is needed

* add guardrails around fields not loaded

* more tests and dataset records cleanup

* refactor a bit more data classes

* more fixes

* refactor dataclass serializers

* more review comments addressed

* add versions setter

* fix falkey listing tests

* add more tests for DatasetListRecord

* cleanup from_dict a bit more

* address review

* cleanup dead code

* more cleanup after reviews

* address review

* mark preview loaded on update dataset

* drop serialized schema from the models

Author: shcheklein

src/datachain/catalog/catalog.py

@@ -5,7 +5,6 @@
 import posixpath
 import sys
 import time
-import traceback
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from contextlib import contextmanager, suppress
 from copy import copy
@@ -66,6 +65,7 @@
     from datachain.data_storage import AbstractMetastore, AbstractWarehouse
     from datachain.dataset import DatasetListVersion
     from datachain.job import Job
+    from datachain.lib.dc.datachain import DataChain
     from datachain.lib.listing_info import ListingInfo
     from datachain.listing import Listing
     from datachain.remote.studio import StudioClient
@@ -77,7 +77,6 @@
 CHECKPOINTS_TTL = 4 * 60 * 60
 
 INDEX_INTERNAL_ERROR_MESSAGE = "Internal error on indexing"
-DATASET_INTERNAL_ERROR_MESSAGE = "Internal error on creating dataset"
 # exit code we use if query script was canceled
 QUERY_SCRIPT_CANCELED_EXIT_CODE = 11
 # exit code we use if the job is already in a terminal state (failed/canceled elsewhere)
@@ -618,17 +617,6 @@ def enlist_source(
 
         return lst, client, list_path
 
-    def _remove_dataset_rows_and_warehouse_info(
-        self, dataset: DatasetRecord, version: str, **kwargs
-    ):
-        self.warehouse.drop_dataset_rows_table(dataset, version)
-        self.update_dataset_version_with_warehouse_info(
-            dataset,
-            version,
-            rows_dropped=True,
-            **kwargs,
-        )
-
     @contextmanager
     def enlist_sources(
         self,
@@ -697,6 +685,7 @@ def _row_to_node(d: dict[str, Any]) -> Node:
                     ds_name,
                     namespace_name=ds_namespace,
                     project_name=ds_project,
+                    versions=[ds_version] if ds_version else None,
                     include_incomplete=False,
                 )
                 if not ds_version:
@@ -802,6 +791,7 @@ def create_dataset(
         columns: Sequence[Column],
         feature_schema: dict | None = None,
         query_script: str = "",
+        sources: str = "",
         validate_version: bool | None = True,
         listing: bool | None = False,
         uuid: str | None = None,
@@ -827,6 +817,7 @@ def create_dataset(
                 name,
                 namespace_name=project.namespace.name if project else None,
                 project_name=project.name if project else None,
+                versions=None,
             )
 
             if (description or attrs) and (
@@ -864,6 +855,7 @@ def create_dataset(
             project=project,
             feature_schema=feature_schema,
             query_script=query_script,
+            sources=sources,
             columns=columns,
             uuid=uuid,
             job_id=job_id,
@@ -873,13 +865,8 @@ def create_dataset(
 
     @staticmethod
     def _next_auto_version(dataset: "DatasetRecord", update_version: str | None) -> str:
-        """Compute the next version for a dataset based on the update strategy.
-
-        Handles brand-new datasets whose versions list may contain a single
-        phantom entry with ``version=None`` (artifact of the LEFT JOIN used
-        by ``get_dataset``).
-        """
-        if not any(v.version for v in dataset.versions):
+        """Compute the next version for a dataset based on the update strategy."""
+        if not dataset.versions:
             return DEFAULT_DATASET_VERSION
         if update_version == "major":
             return dataset.next_version_major
@@ -895,6 +882,7 @@ def _try_claim_version(
         project: Project | None,
         feature_schema: dict | None,
         query_script: str,
+        sources: str,
         columns: Sequence[Column],
         uuid: str | None,
         job_id: str | None,
@@ -929,6 +917,7 @@ def _try_claim_version(
                 target_version,
                 feature_schema=feature_schema,
                 query_script=query_script,
+                sources=sources,
                 columns=columns,
                 uuid=uuid,
                 job_id=job_id,
@@ -953,6 +942,7 @@ def _try_claim_version(
                 name,
                 namespace_name=project.namespace.name if project else None,
                 project_name=project.name if project else None,
+                versions=None,
             )
             target_version = self._next_auto_version(dataset, update_version)
 
@@ -1013,21 +1003,19 @@ def create_dataset_version(
         return dataset, version_created
 
     def update_dataset_version_with_warehouse_info(
-        self, dataset: DatasetRecord, version: str, rows_dropped=False, **kwargs
+        self, dataset: DatasetRecord, version: str, **kwargs
     ) -> None:
         from datachain.query.dataset import DatasetQuery
 
         dataset_version = dataset.get_version(version)
+        if dataset_version._preview_loaded:
+            raise RuntimeError(
+                "update_dataset_version_with_warehouse_info expects preview to be "
+                "unloaded and regenerates it from warehouse rows"
+            )
 
         values = {**kwargs}
 
-        if rows_dropped:
-            values["num_objects"] = None
-            values["size"] = None
-            values["preview"] = None
-            self.metastore.update_dataset_version(dataset, version, **values)
-            return
-
         stats_num_objects = None
         stats_size = None
         if not dataset_version.num_objects:
@@ -1039,22 +1027,20 @@ def update_dataset_version_with_warehouse_info(
             if size != dataset_version.size:
                 values["size"] = size
 
-        preview_rows = None
-        if not dataset_version.preview:
-            preview = (
-                DatasetQuery(
-                    name=dataset.name,
-                    namespace_name=dataset.project.namespace.name,
-                    project_name=dataset.project.name,
-                    version=version,
-                    catalog=self,
-                    include_incomplete=True,  # Allow reading CREATED version
-                )
-                .limit(20)
-                .to_db_records()
+        preview = (
+            DatasetQuery(
+                name=dataset.name,
+                namespace_name=dataset.project.namespace.name,
+                project_name=dataset.project.name,
+                version=version,
+                catalog=self,
+                include_incomplete=True,  # Allow reading CREATED version
             )
-            preview_rows = len(preview)
-            values["preview"] = preview
+            .limit(20)
+            .to_db_records()
+        )
+        preview_rows = len(preview)
+        values["preview"] = preview
 
         # Log anomaly: dataset_stats returned 0 but preview has data
         if stats_num_objects == 0 and preview_rows and preview_rows > 0:
@@ -1068,7 +1054,7 @@ def update_dataset_version_with_warehouse_info(
                 version,
                 dataset_version.num_objects,
                 dataset_version.size,
-                bool(dataset_version.preview),
+                False,
                 stats_num_objects,
                 stats_size,
                 preview_rows,
@@ -1175,7 +1161,7 @@ def create_dataset_from_sources(
         project: Project | None = None,
         client_config=None,
         recursive=False,
-    ) -> DatasetRecord:
+    ) -> "DataChain":
         if not sources:
             raise ValueError("Sources needs to be non empty list")
 
@@ -1188,54 +1174,20 @@ def create_dataset_from_sources(
             if source.startswith(DATASET_PREFIX):
                 dc = read_dataset(source[len(DATASET_PREFIX) :], session=self.session)
             else:
-                dc = read_storage(source, session=self.session, recursive=recursive)
+                dc = read_storage(
+                    source,
+                    session=self.session,
+                    recursive=recursive,
+                    client_config=client_config,
+                )
 
             chains.append(dc)
 
         # create union of all dataset queries created from sources
-        dc = reduce(lambda dc1, dc2: dc1.union(dc2), chains)
-        try:
-            dc = dc.settings(project=project.name, namespace=project.namespace.name)
-            dc.save(name, query_script="")
-        except Exception as e:  # noqa: BLE001
-            try:
-                ds = self.get_dataset(
-                    name,
-                    namespace_name=project.namespace.name,
-                    project_name=project.name,
-                )
-                self.metastore.update_dataset_status(
-                    ds,
-                    DatasetStatus.FAILED,
-                    version=ds.latest_version,
-                    error_message=DATASET_INTERNAL_ERROR_MESSAGE,
-                    error_stack=traceback.format_exc(),
-                )
-                self._remove_dataset_rows_and_warehouse_info(
-                    ds,
-                    ds.latest_version,
-                    sources="\n".join(sources),
-                )
-                raise
-            except DatasetNotFoundError:
-                raise e from None
-
-        ds = self.get_dataset(
-            name,
-            namespace_name=project.namespace.name,
-            project_name=project.name,
-        )
-
-        self.update_dataset_version_with_warehouse_info(
-            ds,
-            ds.latest_version,
-            sources="\n".join(sources),
-        )
-
-        return self.get_dataset(
-            name,
-            namespace_name=project.namespace.name,
-            project_name=project.name,
+        return (
+            reduce(lambda dc1, dc2: dc1.union(dc2), chains)
+            .settings(project=project.name, namespace=project.namespace.name)
+            .save(name, sources="\n".join(sources), query_script="")
         )
 
     def get_full_dataset_name(
@@ -1299,7 +1251,10 @@ def get_dataset(
         name: str,
         namespace_name: str | None = None,
         project_name: str | None = None,
+        *,
+        versions: Sequence[str] | None = (),
         include_incomplete: bool = True,
+        include_preview: bool = False,
     ) -> DatasetRecord:
         from datachain.lib.listing import is_listing_dataset
 
@@ -1314,7 +1269,9 @@ def get_dataset(
             name,
             namespace_name=namespace_name,
             project_name=project_name,
+            versions=versions,
             include_incomplete=include_incomplete,
+            include_preview=include_preview,
         )
 
     def get_dataset_with_remote_fallback(
@@ -1348,6 +1305,7 @@ def get_dataset_with_remote_fallback(
                     name,
                     namespace_name=namespace_name,
                     project_name=project_name,
+                    versions=None,
                     include_incomplete=include_incomplete,
                 )
                 if not version or ds.has_version(version):
@@ -1377,6 +1335,7 @@ def get_dataset_with_remote_fallback(
                 name,
                 namespace_name=namespace_name,
                 project_name=project_name,
+                versions=None,
                 include_incomplete=include_incomplete,
             )
 
@@ -1445,6 +1404,7 @@ def get_dataset_dependencies(
             name,
             namespace_name=namespace_name,
             project_name=project_name,
+            versions=[version],
             include_incomplete=False,
         )
         dataset_version = dataset.get_version(version)
@@ -1618,6 +1578,7 @@ def export_dataset_table(
             name,
             namespace_name=project.namespace.name if project else None,
             project_name=project.name if project else None,
+            versions=[version],
         )
 
         self.warehouse.export_dataset_table(
@@ -1640,6 +1601,7 @@ def remove_dataset(
             name,
             namespace_name=project.namespace.name if project else None,
             project_name=project.name if project else None,
+            versions=None,
         )
         if not version and not force:
             raise ValueError(f"Missing dataset version from input for dataset {name}")
@@ -1680,6 +1642,7 @@ def edit_dataset(
             name,
             namespace_name=project.namespace.name if project else None,
             project_name=project.name if project else None,
+            versions=None,
         )
         return self.update_dataset(dataset, **update_data)
 
@@ -1857,6 +1820,7 @@ def pull_dataset(  # noqa: C901, PLR0915, PLR0912
                     local_ds_name,
                     namespace_name=namespace.name,
                     project_name=project.name,
+                    versions=None,
                     include_incomplete=True,
                 )
                 if local_dataset.has_version(local_ds_version):

src/datachain/cli/commands/datasets.py

@@ -127,6 +127,7 @@ def list_datasets_local_versions(
         name,
         namespace_name=namespace_name,
         project_name=project_name,
+        versions=None,
         include_incomplete=False,
     )
     for v in ds.versions:

src/datachain/cli/commands/show.py

@@ -25,7 +25,11 @@ def show(
         version = name_version
 
     if script:
-        dataset = catalog.get_dataset(name, include_incomplete=False)
+        dataset = catalog.get_dataset(
+            name,
+            versions=[version] if version else None,
+            include_incomplete=False,
+        )
         dataset_version = dataset.get_version(version or dataset.latest_version)
         print(dataset_version.query_script)
         return