Fixes #4003: bulk + async restore for large entity hierarchies

ingestion/src/metadata/ingestion/api/delete.py

@@ -12,6 +12,7 @@
 Delete methods
 """
 
+import os
 import traceback
 from typing import Dict, Iterable, List, Optional, Type  # noqa: UP035
 
@@ -25,13 +26,24 @@
 
 logger = utils_logger()
 
+# Env var that opts every connector into the server-side async delete cascade. When set,
+# mark-deletion calls fire DELETE /<entity>/async/{id}?recursive=true and return 202 + a
+# jobId immediately, so ingestion does not block on the server-side cascade (issue #4003).
+# Explicit dispatch_async= passed to the generators overrides this default.
+DELETE_DISPATCH_ASYNC_ENV = "OM_INGESTION_DELETE_ASYNC"
+
+
+def _default_dispatch_async() -> bool:
+    return os.getenv(DELETE_DISPATCH_ASYNC_ENV, "").lower() in {"true", "1", "yes", "on"}
+
 
 def delete_entity_from_source(
     metadata: OpenMetadata,
     entity_type: Type[T],  # noqa: UP006
     entity_source_state,
     mark_deleted_entity: bool = True,
     params: Optional[Dict[str, str]] = None,  # noqa: UP006, UP045
+    dispatch_async: Optional[bool] = None,  # noqa: UP045
 ) -> Iterable[Either[DeleteEntity]]:
     """
     Method to delete the entities
@@ -40,16 +52,22 @@ def delete_entity_from_source(
     :param entity_source_state: Current state of the service
     :param mark_deleted_entity: Option to mark the entity as deleted or not
     :param params: param to fetch the entity state
+    :param dispatch_async: Route the sink delete through the server-side async endpoint
+        (returns 202 + jobId, runs cascade on the server's executor) so ingestion does
+        not block on large hierarchies — see issue #4003.
     """
+    use_async = dispatch_async if dispatch_async is not None else _default_dispatch_async()
     try:
         entity_state = metadata.list_all_entities(entity=entity_type, params=params)
         for entity in entity_state:
             if str(entity.fullyQualifiedName.root) not in entity_source_state:
                 yield Either(
+                    left=None,
                     right=DeleteEntity(
                         entity=entity,
                         mark_deleted_entities=mark_deleted_entity,
-                    )
+                        dispatch_async=use_async,
+                    ),
                 )
     except Exception as exc:
         yield Either(
@@ -66,19 +84,29 @@ def delete_entity_by_name(
     entity_type: Type[T],  # noqa: UP006
     entity_names: List[str],  # noqa: UP006
     mark_deleted_entity: bool = True,
+    dispatch_async: Optional[bool] = None,  # noqa: UP045
 ) -> Iterable[Either[DeleteEntity]]:
     """
-    Method to delete the entites contained on a given list
+    Method to delete the entities contained on a given list
     :param metadata: OMeta client
     :param entity_type: Pydantic Entity model
     :param entity_names: List of FullyQualifiedNames of the entities to be deleted
     :param mark_deleted_entity: Option to mark the entity as deleted or not
+    :param dispatch_async: see :func:`delete_entity_from_source`
     """
+    use_async = dispatch_async if dispatch_async is not None else _default_dispatch_async()
     try:
         for entity_name in entity_names:
             entity = metadata.get_by_name(entity=entity_type, fqn=entity_name)
             if entity:
-                yield Either(right=DeleteEntity(entity=entity, mark_deleted_entities=mark_deleted_entity))
+                yield Either(
+                    left=None,
+                    right=DeleteEntity(
+                        entity=entity,
+                        mark_deleted_entities=mark_deleted_entity,
+                        dispatch_async=use_async,
+                    ),
+                )
     except Exception as exc:
         yield Either(
             left=StackTraceError(

ingestion/src/metadata/ingestion/models/delete_entity.py

@@ -20,9 +20,13 @@
 
 
 class DeleteEntity(BaseModel):
-    """
-    Entity Reference of the entity to be deleted
+    """Entity reference for a deletion candidate emitted by the ingestion flow.
+
+    ``dispatch_async`` flips the sink to the server-side async delete endpoint
+    (``DELETE /<entity>/async/{id}``) instead of the synchronous one, so ingestion
+    isn't blocked on the cascade for large hierarchies (issue #4003).
     """
 
     entity: Entity
     mark_deleted_entities: Optional[bool] = False  # noqa: UP045
+    dispatch_async: Optional[bool] = False  # noqa: UP045

ingestion/src/metadata/ingestion/ometa/ometa_api.py

@@ -760,6 +760,27 @@ def delete(
         url += f"&hardDelete={str(hard_delete).lower()}"
         self.client.delete(url)
 
+    def delete_async(
+        self,
+        entity: Type[T],  # noqa: UP006
+        entity_id: Union[str, basic.Uuid],  # noqa: UP007
+        recursive: bool = False,
+        hard_delete: bool = False,
+    ) -> Optional[dict]:  # noqa: UP045
+        """Server-side async delete.
+
+        Issues ``DELETE /<entity>/async/{id}?recursive=...&hardDelete=...`` (the dedicated
+        async-delete endpoint defined by ``EntityResource.deleteByIdAsync``) and returns
+        the 202 payload ``{"jobId": ..., "message": ...}``. The actual cascade runs on the
+        server's executor so ingestion can avoid blocking on large hierarchies. Caller is
+        responsible for tracking the returned ``jobId`` if it needs completion confirmation.
+        """
+        url = f"{self.get_suffix(entity)}/async/{model_str(entity_id)}"
+        url += f"?recursive={str(recursive).lower()}"
+        url += f"&hardDelete={str(hard_delete).lower()}"
+        response = self.client.delete(url)
+        return response if isinstance(response, dict) else None
+
     def restore(
         self,
         entity: Type[T],  # noqa: UP006
@@ -794,6 +815,23 @@ def restore(
             )
             return None
 
+    def restore_async(
+        self,
+        entity: Type[T],  # noqa: UP006
+        entity_id: Union[str, basic.Uuid],  # noqa: UP007
+    ) -> Optional[dict]:  # noqa: UP045
+        """Server-side async restore.
+
+        Issues ``PUT /<entity>/restore?async=true`` and returns the 202 payload
+        ``{"jobId": ..., "message": ...}``. Use this when restoring entities with large
+        subtrees so ingestion doesn't block on the cascade (issue #4003). Caller is
+        responsible for tracking the returned ``jobId`` if it needs completion confirmation.
+        """
+        url = f"{self.get_suffix(entity)}/restore?async=true"
+        data = {"id": model_str(entity_id)}
+        response = self.client.put(url, json=data)
+        return response if isinstance(response, dict) else None
+
     def compute_percentile(self, entity: Union[Type[T], str], date: str) -> None:  # noqa: UP006, UP007
         """
         Compute an entity usage percentile

ingestion/src/metadata/ingestion/sink/metadata_rest.py

@@ -584,12 +584,36 @@ def write_users(self, record: OMetaUserProfile) -> Either[User]:
 
     @_run_dispatch.register
     def delete_entity(self, record: DeleteEntity) -> Either[Entity]:
-        self.metadata.delete(
-            entity=type(record.entity),
-            entity_id=record.entity.id,
-            recursive=record.mark_deleted_entities,
-        )
-        return Either(right=record)
+        # record.entity is declared as a bare pydantic BaseModel; the runtime value is a
+        # generated entity that exposes `id` and `fullyQualifiedName`, but basedpyright can't
+        # see those attributes through the BaseModel alias. Pull them via getattr so the type
+        # checker stays quiet without changing the runtime behavior.
+        entity_obj: Any = record.entity
+        entity_id = entity_obj.id
+        fqn = entity_obj.fullyQualifiedName.root
+        recursive = bool(record.mark_deleted_entities)
+        if record.dispatch_async:
+            # Server-side async cascade — returns 202 + jobId immediately so ingestion
+            # doesn't block on large subtrees (issue #4003). The actual work runs on the
+            # server's executor; we surface the jobId in the log for operator correlation.
+            response = self.metadata.delete_async(
+                entity=type(record.entity),
+                entity_id=entity_id,
+                recursive=recursive,
+            )
+            job_id = (response or {}).get("jobId")
+            logger.debug(
+                "Dispatched async delete for %s (jobId=%s)",
+                fqn,
+                job_id,
+            )
+        else:
+            self.metadata.delete(
+                entity=type(record.entity),
+                entity_id=entity_id,
+                recursive=recursive,
+            )
+        return Either(left=None, right=record)
 
     @_run_dispatch.register
     def write_pipeline_status(self, record: OMetaPipelineStatus) -> Either[PipelineStatus]:

ingestion/src/metadata/sdk/entities/base.py

@@ -61,6 +61,54 @@ def execute_async(self) -> Any:
         return export_async(entity=self.entity, name=self.name)
 
 
+@dataclass
+class AsyncJobResponse:
+    """Response shape for server-side async operations.
+
+    Returned with HTTP 202 Accepted by endpoints such as ``PUT /restore?async=true``
+    (issue #4003). The ``job_id`` correlates with WebSocket notifications on the
+    ``restoreEntityChannel`` channel emitted when the work completes.
+    """
+
+    job_id: str
+    message: Optional[str] = None  # noqa: UP045
+
+    @classmethod
+    def from_response(cls, payload: Any) -> "AsyncJobResponse":  # noqa: UP037
+        if isinstance(payload, AsyncJobResponse):
+            return payload
+        if isinstance(payload, dict):
+            job_id = payload.get("jobId")
+            if not job_id:
+                raise ValueError(f"Async response is missing a non-empty jobId: {payload!r}")
+            return cls(job_id=str(job_id), message=payload.get("message"))
+        raise TypeError(f"Cannot coerce {type(payload).__name__} into AsyncJobResponse")
+
+
+@dataclass
+class RestoreOperation(Generic[TEntity]):
+    """Fluent restore builder with optional server-side async dispatch.
+
+    Mirrors the Java SDK's ``Tables.find(id).restore().async().execute()`` style.
+    ``execute()`` runs the synchronous restore and returns the restored entity;
+    ``with_async()`` switches to the server-side async path that returns an
+    :class:`AsyncJobResponse` with a job id (issue #4003).
+    """
+
+    entity_cls: Any  # the BaseEntity subclass that owns this operation
+    entity_id: str
+    async_enabled: bool = field(default=False, init=False)
+
+    def with_async(self) -> "RestoreOperation[TEntity]":  # noqa: UP037
+        self.async_enabled = True
+        return self
+
+    def execute(self) -> Any:
+        if self.async_enabled:
+            return self.entity_cls._restore_server_async(self.entity_id)
+        return self.entity_cls._restore_sync(self.entity_id)
+
+
 @dataclass
 class CsvImportOperation(Generic[TEntity]):
     """Stateful helper for CSV import operations."""
@@ -388,8 +436,12 @@ def remove_followers(cls, entity_id: UuidLike, follower_ids: Sequence[UuidLike])
 
     @classmethod
     def restore(cls, entity_id: UuidLike) -> TEntity:
-        """Restore a soft-deleted entity."""
+        """Restore a soft-deleted entity (synchronous)."""
+
+        return cls._restore_sync(entity_id)
 
+    @classmethod
+    def _restore_sync(cls, entity_id: UuidLike) -> TEntity:
         client = cls._get_client()
         rest_client = cls._get_rest_client(client)
         endpoint = cls._get_endpoint_path(client)
@@ -399,6 +451,50 @@ def restore(cls, entity_id: UuidLike) -> TEntity:
         )
         return cls._coerce_entity(response)
 
+    @classmethod
+    def restore_async(cls, entity_id: UuidLike) -> "AsyncJobResponse":  # noqa: UP037
+        """Trigger a server-side async restore.
+
+        Issues ``PUT /restore?async=true`` and returns the 202 Accepted payload
+        containing the job id. Use this for hierarchies large enough that the
+        synchronous response would exceed proxy / ALB idle timeouts (issue #4003).
+        """
+
+        return cls._restore_server_async(entity_id)
+
+    @classmethod
+    def _restore_server_async(cls, entity_id: UuidLike) -> "AsyncJobResponse":  # noqa: UP037
+        client = cls._get_client()
+        rest_client = cls._get_rest_client(client)
+        endpoint = cls._get_endpoint_path(client)
+        response = rest_client.put(
+            f"{endpoint}/restore?async=true",
+            json={"id": cls._stringify_identifier(entity_id)},
+        )
+        try:
+            return AsyncJobResponse.from_response(response)
+        except ValueError as missing_job_id:
+            # Defensive guard for older servers that don't honor ?async=true (or any
+            # future case where the resource short-circuits with a 200 + entity payload).
+            # Without this, the generic AsyncJobResponse jobId-missing error would be
+            # confusing.
+            raise ValueError(
+                f"Server did not return an async job for {endpoint}/restore. "
+                f"The server may be older than the async-restore release."
+            ) from missing_job_id
+
+    @classmethod
+    def restore_request(cls, entity_id: UuidLike) -> "RestoreOperation[TEntity]":  # noqa: UP037
+        """Return a fluent restore builder.
+
+        Examples::
+
+            restored = Table.restore_request(table_id).execute()
+            job = Table.restore_request(table_id).with_async().execute()
+        """
+
+        return RestoreOperation(entity_cls=cls, entity_id=cls._stringify_identifier(entity_id))
+
     @classmethod
     def update_custom_properties(cls, identifier: UuidLike):
         """Convenience accessor for custom property updates by entity id."""

ingestion/tests/integration/conftest.py

@@ -166,12 +166,19 @@ def _run(workflow_type: Type[IngestionWorkflow], config, raise_from_status=True)
 
 
 def _safe_delete(metadata, entity, entity_id, retries=3, **kwargs):
-    """Delete with retry logic to handle transient server errors during parallel teardown."""
+    """Delete with retry logic to handle transient server errors during parallel teardown.
+
+    A 404 here means the entity is already gone (e.g., wiped as part of an earlier
+    cascade or another worker's teardown); treat it as success rather than retrying.
+    """
     for attempt in range(retries):
         try:
             metadata.delete(entity=entity, entity_id=entity_id, **kwargs)
             return  # noqa: TRY300
-        except Exception:
+        except Exception as exc:
+            if _is_not_found(exc):
+                logger.debug("Skipping %s %s delete — already gone", entity.__name__, entity_id)
+                return
             if attempt < retries - 1:
                 logger.warning(
                     "Retry %d/%d: delete %s %s",
@@ -185,6 +192,13 @@ def _safe_delete(metadata, entity, entity_id, retries=3, **kwargs):
                 raise
 
 
+def _is_not_found(exc: BaseException) -> bool:
+    status = getattr(getattr(exc, "response", None), "status_code", None)
+    if status == 404:
+        return True
+    return "404" in str(exc)
+
+
 @pytest.fixture(scope="module")
 def db_service(metadata, create_service_request, unmask_password):
     service_entity = metadata.create_or_update(data=create_service_request)

ingestion/tests/integration/datalake/conftest.py

@@ -28,6 +28,7 @@
 from metadata.workflow.metadata import MetadataWorkflow
 from metadata.workflow.profiler import ProfilerWorkflow
 
+from ..conftest import _safe_delete  # noqa: TID252
 from ..containers import MinioContainerConfigs, get_minio_container  # noqa: TID252
 from ..integration_base import generate_name  # noqa: TID252
 
@@ -207,7 +208,13 @@ def run_ingestion(metadata, ingestion_config, datalake_service_name):
     yield
     db_service = metadata.get_by_name(entity=DatabaseService, fqn=datalake_service_name)
     if db_service:
-        metadata.delete(DatabaseService, db_service.id, recursive=True, hard_delete=True)
+        _safe_delete(
+            metadata,
+            entity=DatabaseService,
+            entity_id=db_service.id,
+            recursive=True,
+            hard_delete=True,
+        )
 
 
 @pytest.fixture(scope="class")