Ilongin/12872 dataset soft delete

src/datachain/catalog/catalog.py

@@ -48,7 +48,7 @@
     NamespaceNotFoundError,
     ProjectNotFoundError,
 )
-from datachain.lib.listing import get_listing
+from datachain.lib.listing import get_listing, is_listing_dataset
 from datachain.node import DirType, Node, NodeWithPath
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.progress import tqdm
@@ -1079,13 +1079,17 @@ def complete_dataset_version(
         as COMPLETE.
         """
         self.update_dataset_version_with_warehouse_info(dataset, version, **kwargs)
+        # Guard the version-level write: only flip to COMPLETE if the
+        # version is still in a saveable state. Prevents a late-arriving
+        # completion from stomping a concurrent removal.
         self.metastore.update_dataset_status(
             dataset,
             DatasetStatus.COMPLETE,
             version=version,
             error_message=error_message,
             error_stack=error_stack,
             script_output=script_output,
+            where_version_status=[DatasetStatus.CREATED, DatasetStatus.PENDING],
         )
 
     def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
@@ -1095,20 +1099,87 @@ def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
         return dataset_updated
 
     def remove_dataset_version(
-        self, dataset: DatasetRecord, version: str, drop_rows: bool | None = True
+        self, dataset: DatasetRecord, version: str, keep_metadata: bool = True
     ) -> None:
+        """Remove a single dataset version.
+
+        With ``keep_metadata=True`` (default), COMPLETE user-named versions
+        keep a REMOVED record so the semver stays reserved and dependents can
+        resolve lineage; the rows table is always dropped. Non-COMPLETE
+        versions and internal datasets (`lst__*`, `session_*`) are always
+        fully removed. ``keep_metadata=False`` forces a full wipe.
         """
-        Deletes one single dataset version.
-        If it was last version, it removes dataset completely.
-        """
+        from datachain.query.session import Session
+
         if not dataset.has_version(version):
             return
-        self.metastore.update_dataset_version(
-            dataset, version, status=DatasetStatus.REMOVING
+        v = dataset.get_version(version)
+
+        # Resume routing. REMOVING / REMOVING_TOTAL each have a single next
+        # transition, so the current status fully determines the path —
+        # caller's flag is overridden. Internal datasets and incomplete /
+        # failed versions are always wiped.
+        if v.status == DatasetStatus.REMOVING:
+            keep_metadata = True
+        elif (
+            is_listing_dataset(dataset.name)
+            or dataset.name.startswith(Session.DATASET_PREFIX)
+            or v.status
+            in (
+                DatasetStatus.CREATED,
+                DatasetStatus.PENDING,
+                DatasetStatus.FAILED,
+                DatasetStatus.STALE,
+                DatasetStatus.REMOVING_TOTAL,
+            )
+        ):
+            keep_metadata = False
+
+        target = (
+            DatasetStatus.REMOVING if keep_metadata else DatasetStatus.REMOVING_TOTAL
+        )
+        # Atomic claim. Only one caller wins the transition; re-entry from
+        # the same target is allowed for idempotent resume. Wipe is also
+        # allowed from REMOVED so a prior soft-delete can be escalated to a
+        # full wipe.
+        if keep_metadata:
+            allowed_from = [DatasetStatus.COMPLETE, target]
+        else:
+            allowed_from = [
+                DatasetStatus.COMPLETE,
+                DatasetStatus.CREATED,
+                DatasetStatus.PENDING,
+                DatasetStatus.FAILED,
+                DatasetStatus.STALE,
+                DatasetStatus.REMOVED,
+                target,
+            ]
+
+        claimed = self.metastore.update_dataset_version(
+            dataset,
+            version,
+            status=target,
+            where_status=allowed_from,
         )
-        if drop_rows:
+        if claimed is None:
+            logger.debug(
+                "Skipped remove of %s@%s: another caller is already handling it",
+                dataset.name,
+                version,
+            )
+            return
+
+        if keep_metadata:
+            self.warehouse.drop_dataset_rows_table(dataset, version)
+            self.metastore.update_dataset_version(
+                dataset,
+                version,
+                status=DatasetStatus.REMOVED,
+                removed_at=datetime.now(timezone.utc),
+            )
+        else:
             self.warehouse.drop_dataset_rows_table(dataset, version)
-        dataset = self.metastore.remove_dataset_version(dataset, version)
+            self.metastore.remove_dataset_version(dataset, version)
 
     def _remove_versions(self, pairs: Iterable[tuple[DatasetRecord, str]]) -> int:
         num_removed = 0
@@ -1151,7 +1222,7 @@ def cleanup_dataset_versions(self, job_id: str | None = None) -> int:
         Clean up dataset versions that are no longer needed.
 
         Removes dataset versions that:
-        - Have status CREATED, FAILED, STALE, or REMOVING
+        - Have status CREATED, FAILED, STALE, REMOVING, or REMOVING_TOTAL
         - Belong to completed/failed/canceled jobs (not running)
         - Are session_* datasets from finished jobs (orphaned intermediates)
 
@@ -1406,6 +1477,7 @@ def get_dataset_dependencies(
         namespace_name: str | None = None,
         project_name: str | None = None,
         indirect=False,
+        include_removed: bool = True,
     ) -> list[DatasetDependency | None]:
         dataset = self.get_dataset(
             name,
@@ -1422,6 +1494,7 @@ def get_dataset_dependencies(
             return self.metastore.get_direct_dataset_dependencies(
                 dataset,
                 version,
+                include_removed=include_removed,
             )
 
         return self.get_dataset_dependencies_by_ids(
@@ -1603,6 +1676,7 @@ def remove_dataset(
         project: Project | None = None,
         version: str | None = None,
         force: bool | None = False,
+        keep_metadata: bool = True,
     ):
         dataset = self.get_dataset(
             name,
@@ -1618,14 +1692,15 @@ def remove_dataset(
             )
 
         if version:
-            self.remove_dataset_version(dataset, version)
+            self.remove_dataset_version(dataset, version, keep_metadata=keep_metadata)
             return
 
         for v in dataset.versions:
             version = v.version
             self.remove_dataset_version(
                 dataset,
                 version,
+                keep_metadata=keep_metadata,
             )
 
     def edit_dataset(

src/datachain/cli/__init__.py

@@ -196,6 +196,7 @@ def handle_dataset_command(args, catalog):
             args.name,
             version=args.version,
             force=args.force,
+            keep_metadata=args.keep_metadata,
             studio=args.studio,
             team=args.team,
         ),
@@ -204,6 +205,7 @@ def handle_dataset_command(args, catalog):
             args.name,
             version=args.version,
             force=args.force,
+            keep_metadata=args.keep_metadata,
             studio=args.studio,
             team=args.team,
         ),

src/datachain/cli/commands/datasets.py

@@ -153,6 +153,7 @@ def rm_dataset(
     name: str,
     version: str | None = None,
     force: bool | None = False,
+    keep_metadata: bool = True,
     studio: bool | None = False,
     team: str | None = None,
 ) -> None:
@@ -178,7 +179,13 @@ def rm_dataset(
     else:
         try:
             project = catalog.metastore.get_project(project_name, namespace_name)
-            catalog.remove_dataset(name, project, version=version, force=force)
+            catalog.remove_dataset(
+                name,
+                project,
+                version=version,
+                force=force,
+                keep_metadata=keep_metadata,
+            )
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
 

src/datachain/cli/parser/__init__.py

@@ -293,6 +293,17 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         action=BooleanOptionalAction,
         help="Force delete registered dataset with all of its versions",
     )
+    rm_dataset_parser.add_argument(
+        "--keep-metadata",
+        default=True,
+        action=BooleanOptionalAction,
+        help=(
+            "Keep version metadata after removal so the semver stays "
+            "reserved and dependents can still resolve lineage. "
+            "Pass --no-keep-metadata to fully wipe the version "
+            "(rows, dependencies, and version row)."
+        ),
+    )
     rm_dataset_parser.add_argument(
         "--studio",
         action="store_true",