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,14 +1079,25 @@ def complete_dataset_version(
         as COMPLETE.
         """
         self.update_dataset_version_with_warehouse_info(dataset, version, **kwargs)
-        self.metastore.update_dataset_status(
-            dataset,
-            DatasetStatus.COMPLETE,
-            version=version,
-            error_message=error_message,
-            error_stack=error_stack,
-            script_output=script_output,
-        )
+        # 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.
+        try:
+            self.metastore.update_dataset_status(
+                dataset,
+                DatasetStatus.COMPLETE,
+                version=version,
+                error_message=error_message,
+                error_stack=error_stack,
+                script_output=script_output,
+                expected_status=DatasetStatus.CREATED,
+            )
+        except DataChainError as e:
+            raise DataChainError(
+                f"Could not finalize {dataset.name}@{version}: "
+                "the version was removed or modified before save completed. "
+                "This usually means it was deleted concurrently - please retry."
+            ) from e
 
     def update_dataset(self, dataset: DatasetRecord, **kwargs) -> DatasetRecord:
         """Updates dataset fields."""
@@ -1095,20 +1106,144 @@ 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 | None = None,
     ) -> None:
+        """Remove a single dataset version.
+
+        ``keep_metadata=True`` keeps a REMOVED tombstone (the rows table is
+        dropped). ``keep_metadata=False`` fully wipes the version row. When
+        explicit, the call raises if the version's current status is not
+        compatible with the requested mode.
+
+        ``keep_metadata=None`` (default) infers the mode from the version's
+        current status: in-progress removals resume in their existing mode,
+        incomplete or internal-dataset rows are wiped.
         """
-        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
+        v = dataset.get_version(version)
+
+        is_internal = is_listing_dataset(dataset.name) or dataset.name.startswith(
+            Session.DATASET_PREFIX
+        )
+
+        if keep_metadata is None:
+            # Inferred mode (used by GC/cleanup paths): REMOVING resumes soft,
+            # everything else gets wiped.
+            keep_metadata = v.status == DatasetStatus.REMOVING
+        else:
+            # Explicit mode: validate against current status / dataset kind.
+            if keep_metadata and is_internal:
+                raise DataChainError(
+                    f"Internal dataset {dataset.name} cannot be removed "
+                    "while keeping metadata"
+                )
+            if keep_metadata and v.status not in (
+                DatasetStatus.COMPLETE,
+                DatasetStatus.REMOVING,
+                DatasetStatus.REMOVED,
+            ):
+                raise DataChainError(
+                    f"Cannot remove {dataset.name}@{version} while keeping "
+                    f"metadata: current status is {v.status}, expected "
+                    "COMPLETE or REMOVING"
+                )
+            if not keep_metadata and v.status == DatasetStatus.REMOVING:
+                raise DataChainError(
+                    f"Cannot remove {dataset.name}@{version} entirely: "
+                    "a removal that keeps metadata is already in progress"
+                )
+
+        if keep_metadata:
+            self._remove_version_keep_metadata(dataset, version, v.status)
+        else:
+            self._remove_version_wipe_metadata(dataset, version, v.status)
+
+    def _try_claim_transition(
+        self,
+        dataset: DatasetRecord,
+        version: str,
+        *,
+        target_status: int,
+        expected_status: int,
+    ) -> bool:
+        """Atomically transition a version's status; signal win/lost via UUID.
+
+        Two callers matching the same source can both UPDATE; the second one's
+        ``op_uuid`` overwrites the first's. Re-reading after the write reveals
+        which UUID landed - matching ours means we won.
+        """
+        my_uuid = str(uuid4())
+        self.metastore.update_dataset_version(
+            dataset,
+            version,
+            expected_status=expected_status,
+            status=target_status,
+            op_uuid=my_uuid,
+        )
+        try:
+            fresh = self.metastore.get_dataset(
+                dataset.name,
+                namespace_name=dataset.project.namespace.name,
+                project_name=dataset.project.name,
+                versions=[version],
+                include_incomplete=True,
+            )
+            return fresh.get_version(version).op_uuid == my_uuid
+        except (DatasetNotFoundError, DatasetVersionNotFoundError):
+            return False
+
+    def _remove_version_keep_metadata(
+        self, dataset: DatasetRecord, version: str, expected_status: int
+    ) -> None:
+        """Drop rows table, mark version REMOVED (keeps semver + lineage)."""
+        if expected_status == DatasetStatus.REMOVED:
+            return
+        if not self._try_claim_transition(
+            dataset,
+            version,
+            target_status=DatasetStatus.REMOVING,
+            expected_status=expected_status,
+        ):
+            logger.debug(
+                "Skipped remove of %s@%s: another caller is already handling it",
+                dataset.name,
+                version,
+            )
+            return
+
+        self.warehouse.drop_dataset_rows_table(dataset, version)
         self.metastore.update_dataset_version(
-            dataset, version, status=DatasetStatus.REMOVING
+            dataset,
+            version,
+            status=DatasetStatus.REMOVED,
+            removed_at=datetime.now(timezone.utc),
         )
-        if drop_rows:
-            self.warehouse.drop_dataset_rows_table(dataset, version)
-        dataset = self.metastore.remove_dataset_version(dataset, version)
+
+    def _remove_version_wipe_metadata(
+        self, dataset: DatasetRecord, version: str, expected_status: int
+    ) -> None:
+        """Drop rows table and delete the version row entirely."""
+        if not self._try_claim_transition(
+            dataset,
+            version,
+            target_status=DatasetStatus.REMOVING_TOTAL,
+            expected_status=expected_status,
+        ):
+            logger.debug(
+                "Skipped remove of %s@%s: another caller is already handling it",
+                dataset.name,
+                version,
+            )
+            return
+
+        self.warehouse.drop_dataset_rows_table(dataset, version)
+        self.metastore.remove_dataset_version(dataset, version)
 
     def _remove_versions(self, pairs: Iterable[tuple[DatasetRecord, str]]) -> int:
         num_removed = 0
@@ -1151,7 +1286,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 +1541,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 +1558,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 +1740,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 +1756,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",