diff --git a/ingestion/src/metadata/ingestion/api/delete.py b/ingestion/src/metadata/ingestion/api/delete.py index 0653be27c9f8..a1b3ea7b5e4a 100644 --- a/ingestion/src/metadata/ingestion/api/delete.py +++ b/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,6 +26,16 @@ logger = utils_logger() +# Env var that opts every connector into the server-side async delete cascade. When set, +# mark-deletion calls fire DELETE //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, @@ -32,6 +43,7 @@ def delete_entity_from_source( 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( diff --git a/ingestion/src/metadata/ingestion/models/delete_entity.py b/ingestion/src/metadata/ingestion/models/delete_entity.py index f5198db06a8d..7d8dba351d59 100644 --- a/ingestion/src/metadata/ingestion/models/delete_entity.py +++ b/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 //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 diff --git a/ingestion/src/metadata/ingestion/ometa/ometa_api.py b/ingestion/src/metadata/ingestion/ometa/ometa_api.py index 423d3a2a1a58..b92a6e24e0bd 100644 --- a/ingestion/src/metadata/ingestion/ometa/ometa_api.py +++ b/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 //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 //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 diff --git a/ingestion/src/metadata/ingestion/sink/metadata_rest.py b/ingestion/src/metadata/ingestion/sink/metadata_rest.py index 1c6114eac268..5cd5ad5a44b5 100644 --- a/ingestion/src/metadata/ingestion/sink/metadata_rest.py +++ b/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]: diff --git a/ingestion/src/metadata/sdk/entities/base.py b/ingestion/src/metadata/sdk/entities/base.py index 8804e99fc643..f1781d91ad2a 100644 --- a/ingestion/src/metadata/sdk/entities/base.py +++ b/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.""" diff --git a/ingestion/tests/integration/conftest.py b/ingestion/tests/integration/conftest.py index 81a045a58107..e9daf2fd45cb 100644 --- a/ingestion/tests/integration/conftest.py +++ b/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) diff --git a/ingestion/tests/integration/datalake/conftest.py b/ingestion/tests/integration/datalake/conftest.py index c5af9afa1db6..68d73f4a8f57 100644 --- a/ingestion/tests/integration/datalake/conftest.py +++ b/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") diff --git a/ingestion/tests/integration/sdk/conftest.py b/ingestion/tests/integration/sdk/conftest.py index 8136dd3090f8..de9374ad2a42 100644 --- a/ingestion/tests/integration/sdk/conftest.py +++ b/ingestion/tests/integration/sdk/conftest.py @@ -34,6 +34,8 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.workflow.metadata import MetadataWorkflow +from ..conftest import _safe_delete # noqa: TID252 + @pytest.fixture(scope="module") def metadata(): @@ -69,7 +71,13 @@ def db_service(metadata, create_postgres_service, postgres_container): # noqa: service = metadata.get_by_name(DatabaseService, service_entity.fullyQualifiedName.root) if service: - metadata.delete(DatabaseService, service.id, recursive=True, hard_delete=True) + _safe_delete( + metadata, + entity=DatabaseService, + entity_id=service.id, + recursive=True, + hard_delete=True, + ) @pytest.fixture(scope="module") diff --git a/ingestion/tests/unit/sdk/test_restore_async.py b/ingestion/tests/unit/sdk/test_restore_async.py new file mode 100644 index 000000000000..0a014aa23554 --- /dev/null +++ b/ingestion/tests/unit/sdk/test_restore_async.py @@ -0,0 +1,118 @@ +# Copyright 2026 Collate +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the fluent restore + server-side async option (issue #4003).""" + +from unittest.mock import MagicMock + +import pytest + +from metadata.sdk import Tables +from metadata.sdk.entities.base import AsyncJobResponse, RestoreOperation + + +@pytest.fixture +def mock_client(): + client = MagicMock() + client.get_suffix.return_value = "/tables" + rest_client = MagicMock() + client.client = rest_client + Tables.set_default_client(client) + return client + + +def _table_payload(table_id: str) -> dict: + """Minimum dict shape that pydantic_core accepts as a Table.""" + return { + "id": table_id, + "name": "t", + "fullyQualifiedName": "service.db.schema.t", + "deleted": False, + "columns": [], + } + + +def test_restore_sync_calls_put_without_async_param(mock_client): + table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + rest_client = mock_client.client + rest_client.put.return_value = _table_payload(table_id) + + Tables.restore(table_id) + + rest_client.put.assert_called_once() + path = rest_client.put.call_args[0][0] + assert path.endswith("/restore") + assert "async=true" not in path + + +def test_restore_async_appends_async_query_param(mock_client): + table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + rest_client = mock_client.client + rest_client.put.return_value = {"jobId": "job-42", "message": "Restore initiated"} + + response = Tables.restore_async(table_id) + + assert isinstance(response, AsyncJobResponse) + assert response.job_id == "job-42" + assert response.message == "Restore initiated" + rest_client.put.assert_called_once() + path = rest_client.put.call_args[0][0] + assert path.endswith("/restore?async=true") + + +def test_fluent_restore_request_sync_returns_entity(mock_client): + table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + rest_client = mock_client.client + rest_client.put.return_value = _table_payload(table_id) + + op = Tables.restore_request(table_id) + assert isinstance(op, RestoreOperation) + op.execute() + + path = rest_client.put.call_args[0][0] + assert "async=true" not in path + + +def test_fluent_restore_request_with_async_returns_job_response(mock_client): + table_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + rest_client = mock_client.client + rest_client.put.return_value = {"jobId": "job-7", "message": "Restore initiated"} + + job = Tables.restore_request(table_id).with_async().execute() + + assert isinstance(job, AsyncJobResponse) + assert job.job_id == "job-7" + assert "async=true" in rest_client.put.call_args[0][0] + + +def test_async_job_response_from_response_handles_dict(): + response = AsyncJobResponse.from_response({"jobId": "abc", "message": "ok"}) + assert response.job_id == "abc" + assert response.message == "ok" + + +def test_async_job_response_from_response_passes_through_existing(): + original = AsyncJobResponse(job_id="abc", message="ok") + assert AsyncJobResponse.from_response(original) is original + + +def test_async_job_response_from_response_rejects_unknown_type(): + with pytest.raises(TypeError): + AsyncJobResponse.from_response("not a dict") + + +def test_async_job_response_from_response_rejects_missing_job_id(): + with pytest.raises(ValueError, match="non-empty jobId"): + AsyncJobResponse.from_response({"message": "no id here"}) + + +def test_async_job_response_from_response_rejects_empty_job_id(): + with pytest.raises(ValueError, match="non-empty jobId"): + AsyncJobResponse.from_response({"jobId": "", "message": "blank"}) diff --git a/ingestion/tests/unit/test_ometa_restore.py b/ingestion/tests/unit/test_ometa_restore.py index 5b80346ab56a..d50f137c4d33 100644 --- a/ingestion/tests/unit/test_ometa_restore.py +++ b/ingestion/tests/unit/test_ometa_restore.py @@ -131,3 +131,42 @@ def test_restore_endpoint_suffix(self): suffix = metadata.get_suffix(Table) expected_restore_endpoint = f"{suffix}/restore" self.assertEqual(expected_restore_endpoint, "/tables/restore") + + def test_restore_async_dispatches_with_async_query_param(self): + """restore_async should hit /restore?async=true and return the 202 payload.""" + metadata = OpenMetadata(self.server_config) + entity_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + mock_response = {"jobId": "job-42", "message": "Restore initiated successfully."} + + metadata.client.put = MagicMock(return_value=mock_response) + + result = metadata.restore_async(entity=Table, entity_id=entity_id) + + self.assertEqual(result, mock_response) + metadata.client.put.assert_called_once() + call_args = metadata.client.put.call_args + self.assertEqual(call_args[0][0], "/tables/restore?async=true") + self.assertEqual(call_args[1]["json"], {"id": entity_id}) + + def test_delete_async_dispatches_with_async_query_param(self): + """delete_async should hit /async/{id}?recursive=...&hardDelete=... and return the + 202 payload.""" + metadata = OpenMetadata(self.server_config) + entity_id = "b67eac63-9e43-41f5-afb9-387c85df1d8b" + mock_response = {"jobId": "job-7", "message": "Delete initiated successfully."} + + metadata.client.delete = MagicMock(return_value=mock_response) + + result = metadata.delete_async( + entity=Table, + entity_id=entity_id, + recursive=True, + hard_delete=False, + ) + + self.assertEqual(result, mock_response) + metadata.client.delete.assert_called_once() + url = metadata.client.delete.call_args[0][0] + self.assertTrue(url.startswith(f"/tables/async/{entity_id}")) + self.assertIn("recursive=true", url) + self.assertIn("hardDelete=false", url) diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContainerResourceIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContainerResourceIT.java index 000f08003d0e..266444fd77b5 100644 --- a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContainerResourceIT.java +++ b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/ContainerResourceIT.java @@ -1368,14 +1368,15 @@ void test_rootListingExcludesOrphanedChild(TestNamespace ns) { } /** - * Forces the {@code batchDeleteChildren} / {@code processDeletionBatch} path: - * {@code deleteChildren} only takes the batch path when {@code hardDelete=true} AND - * {@code children.size() > 100}. Previously that path pre-deleted relationships in - * two batched queries before iterating {@code cleanup()} per child, and swallowed any - * per-child exception in the loop — so a single failed cleanup left an entity row - * alive with all its relationship rows already wiped (orphan with multi-segment FQN). - * The fix routes everything through {@code cleanup()} per entity and lets exceptions - * propagate. 101 is one above the 100-child threshold that gates the batch path. + * Exercises the {@code bulkHardDeleteSubtree} path that replaced the legacy + * {@code batchDeleteChildren} / {@code processDeletionBatch} flow. The legacy path opened + * an independent JDBI transaction per child via {@code cleanup()} and could leave an + * entity row alive with its relationship rows wiped (orphan with multi-segment FQN) when + * a per-child cleanup failed mid-loop. The replacement runs the entire subtree in a + * single {@code @Transaction} that rolls back atomically on any failure. 101 is one above + * the size that the legacy implementation gated its batch path on — keeping the test + * value pins the regression scenario in place even though the gating threshold no longer + * exists in the code. */ @Test void test_recursiveHardDelete_largeBatch_leavesNoOrphans(TestNamespace ns) { diff --git a/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RestoreHierarchyIT.java b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RestoreHierarchyIT.java new file mode 100644 index 000000000000..94f7e9904796 --- /dev/null +++ b/openmetadata-integration-tests/src/test/java/org/openmetadata/it/tests/RestoreHierarchyIT.java @@ -0,0 +1,321 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.it.tests; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.awaitility.Awaitility; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.openmetadata.it.factories.DatabaseSchemaTestFactory; +import org.openmetadata.it.factories.DatabaseServiceTestFactory; +import org.openmetadata.it.factories.DatabaseTestFactory; +import org.openmetadata.it.factories.GlossaryTermTestFactory; +import org.openmetadata.it.factories.GlossaryTestFactory; +import org.openmetadata.it.factories.TableTestFactory; +import org.openmetadata.it.util.SdkClients; +import org.openmetadata.it.util.TestNamespace; +import org.openmetadata.it.util.TestNamespaceExtension; +import org.openmetadata.schema.entity.data.Database; +import org.openmetadata.schema.entity.data.DatabaseSchema; +import org.openmetadata.schema.entity.data.Glossary; +import org.openmetadata.schema.entity.data.GlossaryTerm; +import org.openmetadata.schema.entity.data.Table; +import org.openmetadata.schema.entity.services.DatabaseService; +import org.openmetadata.schema.type.Include; +import org.openmetadata.sdk.client.OpenMetadataClient; +import org.openmetadata.sdk.fluent.Databases; +import org.openmetadata.sdk.models.AsyncJobResponse; +import org.openmetadata.service.Entity; +import org.openmetadata.service.jdbi3.CollectionDAO; + +/** + * End-to-end tests for the bulk + async restore + bulk hard-delete paths introduced for + * issue #4003 and #4004. + * + *

Builds a small Database → DatabaseSchemas → Tables hierarchy, soft-deletes the database + * (which cascades), then verifies that: + * + *

    + *
  • The synchronous bulk restore path restores the entire subtree in a single PUT call. + *
  • The async restore path returns 202 with a job id and produces the same final state once + * the background work completes. + *
  • The recursive hard-delete on a CONTAINS-shaped service hierarchy wipes every row and + * every entity_relationship reference in one bulk transaction per type. + *
  • The recursive hard-delete on a Glossary → GlossaryTerm hierarchy descends via the + * PARENT_OF relation — confirming the bulk path's relation set covers more than just + * CONTAINS. + *
+ */ +@ExtendWith(TestNamespaceExtension.class) +public class RestoreHierarchyIT { + + private static final int SCHEMA_COUNT = 3; + private static final int TABLES_PER_SCHEMA = 4; + + @BeforeAll + static void setup() { + SdkClients.adminClient(); + } + + @Test + void syncRestore_restoresFullHierarchy(TestNamespace ns) { + Hierarchy h = createHierarchy(ns, "sync"); + softDeleteAndAssertCascade(h); + + Database restored = Databases.find(h.database.getId().toString()).restore().execute(); + assertNotNull(restored); + assertFalse(Boolean.TRUE.equals(restored.getDeleted())); + + assertHierarchyRestored(h); + } + + @Test + void recursiveSoftDelete_marksFullSubtreeDeletedInOnePassPerType(TestNamespace ns) { + Hierarchy h = createHierarchy(ns, "softdel"); + Map recursiveDelete = new HashMap<>(); + recursiveDelete.put("recursive", "true"); + SdkClients.adminClient().databases().delete(h.database.getId().toString(), recursiveDelete); + + OpenMetadataClient client = SdkClients.adminClient(); + Database deletedDb = + client.databases().get(h.database.getId().toString(), "deleted", Include.ALL.value()); + assertTrue(Boolean.TRUE.equals(deletedDb.getDeleted())); + + for (DatabaseSchema schema : h.schemas) { + DatabaseSchema fetched = + client.databaseSchemas().get(schema.getId().toString(), "deleted", Include.ALL.value()); + assertTrue( + Boolean.TRUE.equals(fetched.getDeleted()), + "schema " + schema.getName() + " was not soft-deleted via the bulk cascade"); + } + for (Table table : h.tables) { + Table fetched = client.tables().get(table.getId().toString(), "deleted", Include.ALL.value()); + assertTrue( + Boolean.TRUE.equals(fetched.getDeleted()), + "table " + table.getName() + " was not soft-deleted via the bulk cascade"); + } + } + + @Test + void asyncRestore_returns202AndRestoresFullHierarchy(TestNamespace ns) { + Hierarchy h = createHierarchy(ns, "async"); + softDeleteAndAssertCascade(h); + + AsyncJobResponse job = + Databases.find(h.database.getId().toString()).restore().async().execute(); + assertNotNull(job); + assertNotNull(job.getJobId()); + assertEquals("Restore initiated successfully.", job.getMessage()); + + // Async work runs on the server's executor — poll for completion. + Awaitility.await("async restore for " + h.database.getFullyQualifiedName()) + .atMost(Duration.ofSeconds(60)) + .pollInterval(Duration.ofSeconds(1)) + .ignoreExceptions() + .until( + () -> { + Database current = SdkClients.adminClient().databases().get(h.database.getId()); + return !Boolean.TRUE.equals(current.getDeleted()); + }); + + assertHierarchyRestored(h); + } + + @Test + void hardDelete_databaseService_cascadesEntireSubtreeAndLeavesNoOrphanRelationships( + TestNamespace ns) { + Hierarchy h = createHierarchy(ns, "harddel"); + + Map params = new HashMap<>(); + params.put("recursive", "true"); + params.put("hardDelete", "true"); + SdkClients.adminClient().databaseServices().delete(h.service.getId().toString(), params); + + OpenMetadataClient client = SdkClients.adminClient(); + assertThrows( + Exception.class, + () -> client.databaseServices().get(h.service.getId().toString()), + "database service must be hard-deleted"); + assertThrows( + Exception.class, + () -> client.databases().get(h.database.getId().toString()), + "database must be hard-deleted"); + for (DatabaseSchema schema : h.schemas) { + assertThrows( + Exception.class, + () -> client.databaseSchemas().get(schema.getId().toString()), + "schema must be hard-deleted: " + schema.getName()); + } + for (Table table : h.tables) { + assertThrows( + Exception.class, + () -> client.tables().get(table.getId().toString()), + "table must be hard-deleted: " + table.getName()); + } + + List allDeletedIds = new ArrayList<>(); + allDeletedIds.add(h.service.getId().toString()); + allDeletedIds.add(h.database.getId().toString()); + for (DatabaseSchema schema : h.schemas) { + allDeletedIds.add(schema.getId().toString()); + } + for (Table table : h.tables) { + allDeletedIds.add(table.getId().toString()); + } + assertNoOrphanRelationships(allDeletedIds); + } + + @Test + void hardDelete_glossary_cascadesRecursiveTermsViaParentOf(TestNamespace ns) { + Glossary glossary = GlossaryTestFactory.createWithName(ns, "harddel_glossary"); + GlossaryTerm parent = GlossaryTermTestFactory.createWithName(ns, glossary, "parent_term"); + GlossaryTerm child = GlossaryTermTestFactory.createChild(ns, glossary, parent, "child_term"); + GlossaryTerm grandchild = + GlossaryTermTestFactory.createChild(ns, glossary, child, "grandchild_term"); + + Map params = new HashMap<>(); + params.put("recursive", "true"); + params.put("hardDelete", "true"); + SdkClients.adminClient().glossaries().delete(glossary.getId().toString(), params); + + OpenMetadataClient client = SdkClients.adminClient(); + assertThrows( + Exception.class, + () -> client.glossaries().get(glossary.getId().toString()), + "glossary must be hard-deleted"); + for (GlossaryTerm term : List.of(parent, child, grandchild)) { + assertThrows( + Exception.class, + () -> client.glossaryTerms().get(term.getId().toString()), + "glossary term must be hard-deleted via PARENT_OF cascade: " + term.getName()); + } + + assertNoOrphanRelationships( + List.of( + glossary.getId().toString(), + parent.getId().toString(), + child.getId().toString(), + grandchild.getId().toString())); + } + + private void assertNoOrphanRelationships(List deletedIds) { + CollectionDAO.EntityRelationshipDAO relationshipDAO = + Entity.getCollectionDAO().relationshipDAO(); + List hierarchyRelations = + List.of( + org.openmetadata.schema.type.Relationship.CONTAINS.ordinal(), + org.openmetadata.schema.type.Relationship.PARENT_OF.ordinal(), + org.openmetadata.schema.type.Relationship.HAS.ordinal()); + List outgoing = + relationshipDAO.findToBatchAllTypes(deletedIds, hierarchyRelations, Include.ALL); + assertTrue( + outgoing == null || outgoing.isEmpty(), + "No outgoing entity_relationship rows must reference deleted ids — found " + + (outgoing == null ? 0 : outgoing.size())); + for (Integer relation : hierarchyRelations) { + List incoming = + relationshipDAO.findFromBatch(deletedIds, relation, Include.ALL); + assertTrue( + incoming == null || incoming.isEmpty(), + "No incoming entity_relationship rows must reference deleted ids " + + "(relation=" + + relation + + ") — found " + + (incoming == null ? 0 : incoming.size())); + } + } + + private static class Hierarchy { + DatabaseService service; + Database database; + List schemas; + List tables; + + Hierarchy( + DatabaseService service, + Database database, + List schemas, + List
tables) { + this.service = service; + this.database = database; + this.schemas = schemas; + this.tables = tables; + } + } + + private Hierarchy createHierarchy(TestNamespace ns, String tag) { + DatabaseService service = DatabaseServiceTestFactory.create(ns, "Postgres"); + Database database = DatabaseTestFactory.create(ns, service.getFullyQualifiedName()); + + List schemas = new java.util.ArrayList<>(); + List
tables = new java.util.ArrayList<>(); + for (int s = 0; s < SCHEMA_COUNT; s++) { + DatabaseSchema schema = + DatabaseSchemaTestFactory.create(database.getFullyQualifiedName(), tag + "_schema_" + s); + schemas.add(schema); + for (int t = 0; t < TABLES_PER_SCHEMA; t++) { + tables.add( + TableTestFactory.createSimpleWithName( + tag + "_table_" + s + "_" + t, ns, schema.getFullyQualifiedName())); + } + } + return new Hierarchy(service, database, schemas, tables); + } + + private void softDeleteAndAssertCascade(Hierarchy h) { + Map recursiveDelete = new HashMap<>(); + recursiveDelete.put("recursive", "true"); + SdkClients.adminClient().databases().delete(h.database.getId().toString(), recursiveDelete); + + OpenMetadataClient client = SdkClients.adminClient(); + Database deletedDb = + client.databases().get(h.database.getId().toString(), "deleted", Include.ALL.value()); + assertTrue(Boolean.TRUE.equals(deletedDb.getDeleted()), "database should be soft-deleted"); + + for (DatabaseSchema schema : h.schemas) { + DatabaseSchema fetched = + client.databaseSchemas().get(schema.getId().toString(), "deleted", Include.ALL.value()); + assertTrue(Boolean.TRUE.equals(fetched.getDeleted()), "schema cascade delete failed"); + } + for (Table table : h.tables) { + Table fetched = client.tables().get(table.getId().toString(), "deleted", Include.ALL.value()); + assertTrue(Boolean.TRUE.equals(fetched.getDeleted()), "table cascade delete failed"); + } + } + + private void assertHierarchyRestored(Hierarchy h) { + OpenMetadataClient client = SdkClients.adminClient(); + Database fetchedDb = client.databases().get(h.database.getId().toString()); + assertFalse(Boolean.TRUE.equals(fetchedDb.getDeleted()), "database not restored"); + + for (DatabaseSchema schema : h.schemas) { + DatabaseSchema fetched = client.databaseSchemas().get(schema.getId().toString()); + assertFalse(Boolean.TRUE.equals(fetched.getDeleted()), "schema not restored"); + } + for (Table table : h.tables) { + Table fetched = client.tables().get(table.getId().toString()); + assertFalse(Boolean.TRUE.equals(fetched.getDeleted()), "table not restored"); + } + } +} diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIApplications.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIApplications.java index 620094ebdbbc..d783ad381740 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIApplications.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIApplications.java @@ -213,6 +213,11 @@ public AIApplication get() { public AIApplicationDeleter delete() { return new AIApplicationDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.aiApplications(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIGovernancePolicies.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIGovernancePolicies.java index 678f941cdbc6..12fff1ca9902 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIGovernancePolicies.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/AIGovernancePolicies.java @@ -215,6 +215,11 @@ public AIGovernancePolicy get() { public AIGovernancePolicyDeleter delete() { return new AIGovernancePolicyDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.aiGovernancePolicies(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Charts.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Charts.java index eeda8f4ad421..b726153be885 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Charts.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Charts.java @@ -177,6 +177,10 @@ public FluentChart fetch() { public ChartDeleter delete() { return new ChartDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.charts(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Classifications.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Classifications.java index 7e9ab3eb36b9..c07463ea9a3f 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Classifications.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Classifications.java @@ -175,6 +175,11 @@ public FluentClassification fetch() { public ClassificationDeleter delete() { return new ClassificationDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.classifications(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Containers.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Containers.java index a2f847834a2c..ca35c0341ef6 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Containers.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Containers.java @@ -235,6 +235,11 @@ public FluentContainer fetch() { public ContainerDeleter delete() { return new ContainerDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.containers(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DashboardDataModels.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DashboardDataModels.java index 8f109552a41a..3f2dfa65557b 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DashboardDataModels.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DashboardDataModels.java @@ -192,6 +192,11 @@ public FluentDashboardDataModel fetch() { public DashboardDataModelDeleter delete() { return new DashboardDataModelDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.dashboardDataModels(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Dashboards.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Dashboards.java index 4af581aadef5..ed2585a61e31 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Dashboards.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Dashboards.java @@ -183,6 +183,11 @@ public FluentDashboard fetch() { public DashboardDeleter delete() { return new DashboardDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.dashboards(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataContracts.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataContracts.java index be6a50955f36..a7b248437207 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataContracts.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataContracts.java @@ -261,6 +261,11 @@ public FluentDataContract fetch() { public DataContractDeleter delete() { return new DataContractDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.dataContracts(), identifier); + } } // ==================== Contract Operations ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataProducts.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataProducts.java index 5b912ba37487..e932c745a613 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataProducts.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DataProducts.java @@ -180,6 +180,11 @@ public FluentDataProduct fetch() { public DataProductDeleter delete() { return new DataProductDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.dataProducts(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DatabaseSchemas.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DatabaseSchemas.java index bb8e03771e59..948b6786ba37 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DatabaseSchemas.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/DatabaseSchemas.java @@ -248,6 +248,11 @@ public FluentDatabaseSchema fetch() { public DatabaseSchemaDeleter delete() { return new DatabaseSchemaDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.databaseSchemas(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Databases.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Databases.java index 7741c54c1f07..91b6a3046563 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Databases.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Databases.java @@ -246,6 +246,11 @@ public FluentDatabase fetch() { public DatabaseDeleter delete() { return new DatabaseDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.databases(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Domains.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Domains.java index b5d3d76404a7..b28d61dc20e6 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Domains.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Domains.java @@ -223,6 +223,10 @@ public FluentDomain fetch() { public DomainDeleter delete() { return new DomainDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.domains(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Glossaries.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Glossaries.java index ffac0694634a..6e0868e15e50 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Glossaries.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Glossaries.java @@ -193,6 +193,11 @@ public FluentGlossary fetch() { public GlossaryDeleter delete() { return new GlossaryDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.glossaries(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/GlossaryTerms.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/GlossaryTerms.java index 6370aa78eb60..fc31477a9a23 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/GlossaryTerms.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/GlossaryTerms.java @@ -227,6 +227,11 @@ public FluentGlossaryTerm fetch() { public GlossaryTermDeleter delete() { return new GlossaryTermDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.glossaryTerms(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/LLMServices.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/LLMServices.java index f8b0cebc895f..25efdd7beaea 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/LLMServices.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/LLMServices.java @@ -214,6 +214,11 @@ public LLMService get() { public LLMServiceDeleter delete() { return new LLMServiceDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.llmServices(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/McpServers.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/McpServers.java index f3b88c399ef8..dd982f606c6c 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/McpServers.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/McpServers.java @@ -290,6 +290,11 @@ public McpServer get() { public McpServerDeleter delete() { return new McpServerDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.mcpServers(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Metrics.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Metrics.java index 36f96fbb7fd8..6b2ed1bd09cf 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Metrics.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Metrics.java @@ -182,6 +182,10 @@ public FluentMetric fetch() { public MetricDeleter delete() { return new MetricDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.metrics(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/MlModels.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/MlModels.java index 263270d9dc32..6dd864075f68 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/MlModels.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/MlModels.java @@ -178,6 +178,10 @@ public FluentMlModel fetch() { public MlModelDeleter delete() { return new MlModelDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.mlModels(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Pipelines.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Pipelines.java index 8097fa8cc4d3..fe8d55502bc0 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Pipelines.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Pipelines.java @@ -253,6 +253,11 @@ public FluentPipeline fetch() { public PipelineDeleter delete() { return new PipelineDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.pipelines(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/PromptTemplates.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/PromptTemplates.java index 7af80a23323f..4670a75b23e0 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/PromptTemplates.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/PromptTemplates.java @@ -207,6 +207,11 @@ public PromptTemplate get() { public PromptTemplateDeleter delete() { return new PromptTemplateDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.promptTemplates(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Queries.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Queries.java index 0b300bd7f4b9..35f464944b53 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Queries.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Queries.java @@ -177,6 +177,10 @@ public FluentQuery fetch() { public QueryDeleter delete() { return new QueryDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.queries(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/SearchIndexes.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/SearchIndexes.java index ec810d5cdfd5..c37e07fd57c3 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/SearchIndexes.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/SearchIndexes.java @@ -180,6 +180,11 @@ public FluentSearchIndex fetch() { public SearchIndexDeleter delete() { return new SearchIndexDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.searchIndexes(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/StoredProcedures.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/StoredProcedures.java index 8ef24a22f032..d808adb73960 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/StoredProcedures.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/StoredProcedures.java @@ -190,6 +190,11 @@ public FluentStoredProcedure fetch() { public StoredProcedureDeleter delete() { return new StoredProcedureDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>( + client.storedProcedures(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tables.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tables.java index c11760f1f2d9..f63d4c86bbbc 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tables.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tables.java @@ -312,6 +312,10 @@ public org.openmetadata.sdk.fluent.wrappers.FluentTable fetch() { public TableDeleter delete() { return new TableDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer
restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.tables(), identifier); + } } // ==================== Table Operations ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tags.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tags.java index dd0ac56e8142..5bdd15b2aca8 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tags.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Tags.java @@ -177,6 +177,10 @@ public FluentTag fetch() { public TagDeleter delete() { return new TagDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.tags(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Teams.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Teams.java index 95e7de4cee63..43db1de97327 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Teams.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Teams.java @@ -233,6 +233,10 @@ public FluentTeam fetch() { public TeamDeleter delete() { return new TeamDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.teams(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Topics.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Topics.java index f31328e0b714..09bfe323eed2 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Topics.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Topics.java @@ -238,6 +238,10 @@ public FluentTopic fetch() { public TopicDeleter delete() { return new TopicDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.topics(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Users.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Users.java index 35d54df0bb27..0c9c49fab84b 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Users.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/Users.java @@ -267,6 +267,10 @@ public FluentUser fetch() { public UserDeleter delete() { return new UserDeleter(client, identifier); } + + public org.openmetadata.sdk.fluent.common.EntityRestorer restore() { + return new org.openmetadata.sdk.fluent.common.EntityRestorer<>(client.users(), identifier); + } } // ==================== Deleter ==================== diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/AsyncEntityRestorer.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/AsyncEntityRestorer.java new file mode 100644 index 000000000000..6419d7a38b92 --- /dev/null +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/AsyncEntityRestorer.java @@ -0,0 +1,39 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.sdk.fluent.common; + +import org.openmetadata.sdk.models.AsyncJobResponse; +import org.openmetadata.sdk.services.EntityServiceBase; + +/** + * Generic fluent async restore builder. Returned by {@link EntityRestorer#async()}. + * Calls {@link EntityServiceBase#restoreServerAsync(String)} which issues + * {@code PUT /restore?async=true} and returns the 202 Accepted response carrying the + * job id (issue #4003). The {@code } parameter is preserved for symmetry with + * {@link EntityRestorer} so call sites that already have an + * {@code EntityRestorer} reference can switch to the async variant without + * losing the type-level context, even though the response itself is type-erased. + */ +public class AsyncEntityRestorer { + private final EntityServiceBase service; + private final String id; + + public AsyncEntityRestorer(EntityServiceBase service, String id) { + this.service = service; + this.id = id; + } + + public AsyncJobResponse execute() { + return service.restoreServerAsync(id); + } +} diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/EntityRestorer.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/EntityRestorer.java new file mode 100644 index 000000000000..11b7573605fb --- /dev/null +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/fluent/common/EntityRestorer.java @@ -0,0 +1,45 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.sdk.fluent.common; + +import org.openmetadata.sdk.services.EntityServiceBase; + +/** + * Generic fluent restore builder used by every entity-type fluent class that exposes a + * {@code restore()} entry point (Tables, Dashboards, Pipelines, Topics, Containers, + * Glossaries, Domains, …). Replaces the per-entity {@code TableRestorer} / + * {@code DatabaseRestorer} duplicates so adding restore support to a new fluent only + * requires wiring it to its service — no new class per type. + * + *

Sync: {@code execute()} runs the synchronous restore and returns the restored + * entity. Async: {@code async().execute()} switches to the server-side async path + * ({@code PUT /restore?async=true}) and returns an + * {@link org.openmetadata.sdk.models.AsyncJobResponse} with a job id (issue #4003). + */ +public class EntityRestorer { + private final EntityServiceBase service; + private final String id; + + public EntityRestorer(EntityServiceBase service, String id) { + this.service = service; + this.id = id; + } + + public AsyncEntityRestorer async() { + return new AsyncEntityRestorer<>(service, id); + } + + public T execute() { + return service.restore(id); + } +} diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/models/AsyncJobResponse.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/models/AsyncJobResponse.java new file mode 100644 index 000000000000..fd39a1bc4ecf --- /dev/null +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/models/AsyncJobResponse.java @@ -0,0 +1,49 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.sdk.models; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; + +/** + * Response shape returned with HTTP 202 Accepted for server-side async operations such as + * restore (issue #4003) and delete. Contains a job id that can be used with the SDK + * WebSocketListener to await completion notifications. + */ +@JsonIgnoreProperties(ignoreUnknown = true) +public class AsyncJobResponse { + private String jobId; + private String message; + + public AsyncJobResponse() {} + + public AsyncJobResponse(String jobId, String message) { + this.jobId = jobId; + this.message = message; + } + + public String getJobId() { + return jobId; + } + + public void setJobId(String jobId) { + this.jobId = jobId; + } + + public String getMessage() { + return message; + } + + public void setMessage(String message) { + this.message = message; + } +} diff --git a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/services/EntityServiceBase.java b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/services/EntityServiceBase.java index 14f5f7d72259..1a7a2be3b352 100644 --- a/openmetadata-sdk/src/main/java/org/openmetadata/sdk/services/EntityServiceBase.java +++ b/openmetadata-sdk/src/main/java/org/openmetadata/sdk/services/EntityServiceBase.java @@ -421,15 +421,17 @@ public T restore(String id) throws OpenMetadataException { } /** - * Restore a soft-deleted entity (async) + * Restore a soft-deleted entity (client-side async wrapper). + * + *

Runs the synchronous restore call on the SDK's executor and returns a + * {@link CompletableFuture}. The server still does the work synchronously inside the request, + * so this still ties up an HTTP connection for the duration. For large hierarchies use + * {@link #restoreServerAsync(String)} instead, which returns a 202 with a job id. */ public CompletableFuture restoreAsync(UUID id) { return restoreAsync(id.toString()); } - /** - * Restore a soft-deleted entity (async) - */ public CompletableFuture restoreAsync(String id) { org.openmetadata.schema.api.data.RestoreEntity restoreEntity = new org.openmetadata.schema.api.data.RestoreEntity(); @@ -438,6 +440,43 @@ public CompletableFuture restoreAsync(String id) { HttpMethod.PUT, basePath + "/restore", restoreEntity, getEntityClass()); } + /** + * Trigger a server-side async restore. Issues {@code PUT /restore?async=true} and returns + * the 202 Accepted response containing the job id. Used to avoid proxy / ALB idle timeouts + * on large hierarchies (issue #4003). The caller can await completion via the SDK's + * WebSocketListener on the {@code restoreEntityChannel} channel. + */ + public org.openmetadata.sdk.models.AsyncJobResponse restoreServerAsync(UUID id) + throws OpenMetadataException { + return restoreServerAsync(id.toString()); + } + + public org.openmetadata.sdk.models.AsyncJobResponse restoreServerAsync(String id) + throws OpenMetadataException { + org.openmetadata.schema.api.data.RestoreEntity restoreEntity = + new org.openmetadata.schema.api.data.RestoreEntity(); + restoreEntity.setId(java.util.UUID.fromString(id)); + RequestOptions options = RequestOptions.builder().queryParam("async", "true").build(); + org.openmetadata.sdk.models.AsyncJobResponse response = + httpClient.execute( + HttpMethod.PUT, + basePath + "/restore", + restoreEntity, + org.openmetadata.sdk.models.AsyncJobResponse.class, + options); + // Defensive check for older servers that don't honor ?async=true (or for any future + // case where the resource short-circuits with a 200 + entity payload). Jackson would + // otherwise silently deserialize the entity JSON into an AsyncJobResponse with all + // null fields and callers would treat a sync restore as a dispatched async job. + if (response == null || response.getJobId() == null || response.getJobId().isEmpty()) { + throw new OpenMetadataException( + "Server did not return an async job for " + + basePath + + "/restore. The server may be older than the async-restore release."); + } + return response; + } + /** * Export entity data to CSV format. * diff --git a/openmetadata-sdk/src/test/java/org/openmetadata/sdk/fluent/RestoreFluentAPITest.java b/openmetadata-sdk/src/test/java/org/openmetadata/sdk/fluent/RestoreFluentAPITest.java new file mode 100644 index 000000000000..23ad3860d1df --- /dev/null +++ b/openmetadata-sdk/src/test/java/org/openmetadata/sdk/fluent/RestoreFluentAPITest.java @@ -0,0 +1,270 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.sdk.fluent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.UUID; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.openmetadata.schema.entity.data.Container; +import org.openmetadata.schema.entity.data.Dashboard; +import org.openmetadata.schema.entity.data.Database; +import org.openmetadata.schema.entity.data.DatabaseSchema; +import org.openmetadata.schema.entity.data.Glossary; +import org.openmetadata.schema.entity.data.MlModel; +import org.openmetadata.schema.entity.data.Pipeline; +import org.openmetadata.schema.entity.data.Table; +import org.openmetadata.schema.entity.data.Topic; +import org.openmetadata.schema.entity.domains.Domain; +import org.openmetadata.sdk.client.OpenMetadataClient; +import org.openmetadata.sdk.models.AsyncJobResponse; +import org.openmetadata.sdk.services.dataassets.DashboardService; +import org.openmetadata.sdk.services.dataassets.MlModelService; +import org.openmetadata.sdk.services.dataassets.PipelineService; +import org.openmetadata.sdk.services.dataassets.TableService; +import org.openmetadata.sdk.services.dataassets.TopicService; +import org.openmetadata.sdk.services.databases.DatabaseSchemaService; +import org.openmetadata.sdk.services.databases.DatabaseService; +import org.openmetadata.sdk.services.domains.DomainService; +import org.openmetadata.sdk.services.glossary.GlossaryService; +import org.openmetadata.sdk.services.storages.ContainerService; + +/** + * Verifies the fluent restore builders added for issue #4003. {@code .restore().execute()} + * routes to the synchronous SDK call; chaining {@code .async()} switches to the server-side + * async path and returns an {@link AsyncJobResponse}. + */ +class RestoreFluentAPITest { + + @Mock private OpenMetadataClient mockClient; + @Mock private TableService mockTables; + @Mock private DatabaseService mockDatabases; + @Mock private DatabaseSchemaService mockSchemas; + @Mock private DashboardService mockDashboards; + @Mock private PipelineService mockPipelines; + @Mock private TopicService mockTopics; + @Mock private MlModelService mockMlModels; + @Mock private ContainerService mockContainers; + @Mock private GlossaryService mockGlossaries; + @Mock private DomainService mockDomains; + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + when(mockClient.tables()).thenReturn(mockTables); + when(mockClient.databases()).thenReturn(mockDatabases); + when(mockClient.databaseSchemas()).thenReturn(mockSchemas); + when(mockClient.dashboards()).thenReturn(mockDashboards); + when(mockClient.pipelines()).thenReturn(mockPipelines); + when(mockClient.topics()).thenReturn(mockTopics); + when(mockClient.mlModels()).thenReturn(mockMlModels); + when(mockClient.containers()).thenReturn(mockContainers); + when(mockClient.glossaries()).thenReturn(mockGlossaries); + when(mockClient.domains()).thenReturn(mockDomains); + Tables.setDefaultClient(mockClient); + Databases.setDefaultClient(mockClient); + DatabaseSchemas.setDefaultClient(mockClient); + Dashboards.setDefaultClient(mockClient); + Pipelines.setDefaultClient(mockClient); + Topics.setDefaultClient(mockClient); + MlModels.setDefaultClient(mockClient); + Containers.setDefaultClient(mockClient); + Glossaries.setDefaultClient(mockClient); + Domains.setDefaultClient(mockClient); + } + + @Test + void tablesFluent_syncRestore_callsRestore() throws Exception { + String id = UUID.randomUUID().toString(); + Table restored = new Table().withId(UUID.fromString(id)).withName("t"); + when(mockTables.restore(id)).thenReturn(restored); + + Table result = Tables.find(id).restore().execute(); + + assertSame(restored, result); + verify(mockTables).restore(id); + verify(mockTables, never()).restoreServerAsync(eq(id)); + } + + @Test + void tablesFluent_asyncRestore_callsRestoreServerAsync() throws Exception { + String id = UUID.randomUUID().toString(); + AsyncJobResponse expected = new AsyncJobResponse("job-1", "Restore initiated successfully."); + when(mockTables.restoreServerAsync(id)).thenReturn(expected); + + AsyncJobResponse result = Tables.find(id).restore().async().execute(); + + assertNotNull(result); + assertEquals("job-1", result.getJobId()); + assertEquals("Restore initiated successfully.", result.getMessage()); + verify(mockTables).restoreServerAsync(id); + verify(mockTables, never()).restore(eq(id)); + } + + @Test + void databasesFluent_syncRestore_callsRestore() throws Exception { + String id = UUID.randomUUID().toString(); + Database restored = new Database().withId(UUID.fromString(id)).withName("db"); + when(mockDatabases.restore(id)).thenReturn(restored); + + Database result = Databases.find(id).restore().execute(); + + assertSame(restored, result); + verify(mockDatabases).restore(id); + verify(mockDatabases, never()).restoreServerAsync(eq(id)); + } + + @Test + void databasesFluent_asyncRestore_callsRestoreServerAsync() throws Exception { + String id = UUID.randomUUID().toString(); + AsyncJobResponse expected = new AsyncJobResponse("job-2", "Restore initiated successfully."); + when(mockDatabases.restoreServerAsync(id)).thenReturn(expected); + + AsyncJobResponse result = Databases.find(id).restore().async().execute(); + + assertNotNull(result); + assertEquals("job-2", result.getJobId()); + verify(mockDatabases).restoreServerAsync(id); + verify(mockDatabases, never()).restore(eq(id)); + } + + // ---------------------------------------------------------------------------------------- + // Coverage that the new generic EntityRestorer wiring works for every data-asset fluent. + // Tables / Databases above are unchanged; below verifies the broader rollout reaches the + // correct service per fluent — one sync + one async assertion per type to lock the + // wiring without exhaustively testing every type (they all go through the same + // EntityRestorer helper, so a representative sample is enough to catch a typo in any + // single fluent's wire-up). + // ---------------------------------------------------------------------------------------- + + @Test + void databaseSchemasFluent_restore_routesThroughSchemaService() throws Exception { + String id = UUID.randomUUID().toString(); + DatabaseSchema restored = new DatabaseSchema().withId(UUID.fromString(id)).withName("s"); + when(mockSchemas.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-schema", "ok"); + when(mockSchemas.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, DatabaseSchemas.find(id).restore().execute()); + assertEquals("job-schema", DatabaseSchemas.find(id).restore().async().execute().getJobId()); + verify(mockSchemas).restore(id); + verify(mockSchemas).restoreServerAsync(id); + } + + @Test + void dashboardsFluent_restore_routesThroughDashboardService() throws Exception { + String id = UUID.randomUUID().toString(); + Dashboard restored = new Dashboard().withId(UUID.fromString(id)).withName("d"); + when(mockDashboards.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-dash", "ok"); + when(mockDashboards.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Dashboards.find(id).restore().execute()); + assertEquals("job-dash", Dashboards.find(id).restore().async().execute().getJobId()); + verify(mockDashboards).restore(id); + verify(mockDashboards).restoreServerAsync(id); + } + + @Test + void pipelinesFluent_restore_routesThroughPipelineService() throws Exception { + String id = UUID.randomUUID().toString(); + Pipeline restored = new Pipeline().withId(UUID.fromString(id)).withName("p"); + when(mockPipelines.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-pipe", "ok"); + when(mockPipelines.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Pipelines.find(id).restore().execute()); + assertEquals("job-pipe", Pipelines.find(id).restore().async().execute().getJobId()); + verify(mockPipelines).restore(id); + verify(mockPipelines).restoreServerAsync(id); + } + + @Test + void topicsFluent_restore_routesThroughTopicService() throws Exception { + String id = UUID.randomUUID().toString(); + Topic restored = new Topic().withId(UUID.fromString(id)).withName("t"); + when(mockTopics.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-topic", "ok"); + when(mockTopics.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Topics.find(id).restore().execute()); + assertEquals("job-topic", Topics.find(id).restore().async().execute().getJobId()); + verify(mockTopics).restore(id); + verify(mockTopics).restoreServerAsync(id); + } + + @Test + void mlModelsFluent_restore_routesThroughMlModelService() throws Exception { + String id = UUID.randomUUID().toString(); + MlModel restored = new MlModel().withId(UUID.fromString(id)).withName("m"); + when(mockMlModels.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-ml", "ok"); + when(mockMlModels.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, MlModels.find(id).restore().execute()); + assertEquals("job-ml", MlModels.find(id).restore().async().execute().getJobId()); + verify(mockMlModels).restore(id); + verify(mockMlModels).restoreServerAsync(id); + } + + @Test + void containersFluent_restore_routesThroughContainerService() throws Exception { + String id = UUID.randomUUID().toString(); + Container restored = new Container().withId(UUID.fromString(id)).withName("c"); + when(mockContainers.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-cont", "ok"); + when(mockContainers.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Containers.find(id).restore().execute()); + assertEquals("job-cont", Containers.find(id).restore().async().execute().getJobId()); + verify(mockContainers).restore(id); + verify(mockContainers).restoreServerAsync(id); + } + + @Test + void glossariesFluent_restore_routesThroughGlossaryService() throws Exception { + String id = UUID.randomUUID().toString(); + Glossary restored = new Glossary().withId(UUID.fromString(id)).withName("g"); + when(mockGlossaries.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-gloss", "ok"); + when(mockGlossaries.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Glossaries.find(id).restore().execute()); + assertEquals("job-gloss", Glossaries.find(id).restore().async().execute().getJobId()); + verify(mockGlossaries).restore(id); + verify(mockGlossaries).restoreServerAsync(id); + } + + @Test + void domainsFluent_restore_routesThroughDomainService() throws Exception { + String id = UUID.randomUUID().toString(); + Domain restored = new Domain().withId(UUID.fromString(id)).withName("dom"); + when(mockDomains.restore(id)).thenReturn(restored); + AsyncJobResponse async = new AsyncJobResponse("job-dom", "ok"); + when(mockDomains.restoreServerAsync(id)).thenReturn(async); + + assertSame(restored, Domains.find(id).restore().execute()); + assertEquals("job-dom", Domains.find(id).restore().async().execute().getJobId()); + verify(mockDomains).restore(id); + verify(mockDomains).restoreServerAsync(id); + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java index 0f08e3c21f99..af4d16cb744d 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java @@ -2130,6 +2130,29 @@ default List findToBatchAllTypes( return findToBatchAllTypesWithCondition(fromIds, relation, condition); } + @SqlQuery( + "SELECT fromId, toId, fromEntity, toEntity, relation, json, jsonSchema " + + "FROM entity_relationship " + + "WHERE fromId IN () " + + "AND relation IN () " + + "") + @UseRowMapper(RelationshipObjectMapper.class) + List findToBatchAllTypesWithRelationsCondition( + @BindList("fromIds") List fromIds, + @BindList("relations") List relations, + @Define("cond") String condition); + + default List findToBatchAllTypes( + List fromIds, List relations, Include include) { + String condition = ""; + if (include == null || include == Include.NON_DELETED) { + condition = "AND deleted = FALSE"; + } else if (include == Include.DELETED) { + condition = "AND deleted = TRUE"; + } + return findToBatchAllTypesWithRelationsCondition(fromIds, relations, condition); + } + @SqlQuery( "SELECT fromId, toId, fromEntity, toEntity, relation, json, jsonSchema " + "FROM entity_relationship " @@ -4171,6 +4194,16 @@ List findByEntityId( @SqlQuery("select id from thread_entity where entityId = :entityId") List findByEntityId(@Bind("entityId") String entityId); + // DISTINCT is defence-in-depth: thread_entity.id is a primary key, and entityId is a + // single-valued column per row, so a single matching scan can't physically return the + // same id twice. The DISTINCT survives a future schema where a thread row picks up + // multiple entity references (or a join is added) — keeping the consumer code in + // deleteByAbout from re-issuing redundant relationship / extension / feed deletes for + // the same id under chunking. + @SqlQuery("select DISTINCT id from where entityId IN ()") + List findByEntityIds( + @Define("tableName") String tableName, @BindList("entityIds") List entityIds); + @ConnectionAwareSqlUpdate( value = "UPDATE SET json = JSON_SET(json, '$.about', :newEntityLink)\n" diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardRepository.java index d7a3d0507f9e..371f155a3296 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/DashboardRepository.java @@ -209,14 +209,28 @@ public void clearFields(Dashboard dashboard, Fields fields) { fields.contains("usageSummary") ? dashboard.getUsageSummary() : null); } - // Override soft delete behavior to handle charts through HAS relation. + // Hard-delete chart links (HAS relation). The CONTAINS subtree is handled by the bulk + // path in EntityRepository.bulkHardDeleteSubtree; chart handling is a per-dashboard concern + // and lives in the per-entity extension hook so it runs both for direct dashboard deletes + // and when dashboards are descendants of a larger hard-delete cascade. @Transaction @Override - protected void deleteChildren( - UUID dashboardId, boolean recursive, boolean hardDelete, String updatedBy) { - super.deleteChildren(dashboardId, recursive, hardDelete, updatedBy); + protected void hardDeleteAdditionalChildren(UUID dashboardId, String updatedBy) { + cascadeChartCleanup(dashboardId, updatedBy, true); + } - // Load all charts linked to this dashboard + // Soft-delete chart links (HAS relation). The CONTAINS subtree is handled by the bulk + // path in EntityRepository.bulkSoftDeleteSubtree; chart handling is a per-dashboard + // concern and lives in the per-entity extension hook so it runs both for direct dashboard + // deletes and when dashboards are descendants of a larger soft-delete (e.g., + // DashboardService cascade). + @Transaction + @Override + protected void softDeleteAdditionalChildren(UUID dashboardId, String updatedBy) { + cascadeChartCleanup(dashboardId, updatedBy, false); + } + + private void cascadeChartCleanup(UUID dashboardId, String updatedBy, boolean hardDelete) { List chartRecords = daoCollection .relationshipDAO() @@ -225,7 +239,6 @@ protected void deleteChildren( return; } - // Batch-load dashboard relationships for these charts List dashboardRelationships = daoCollection .relationshipDAO() @@ -248,11 +261,10 @@ protected void deleteChildren( Include.NON_DELETED) .stream() .map(Dashboard::getId) - .filter(id -> !id.equals(dashboardId)) // (excluding the current dashboard + .filter(id -> !id.equals(dashboardId)) .collect(Collectors.toSet()); - // For deletion: get charts whose linked dashboards (excluding the current dashboard) - // have no other non‑deleted dashboards. + // Soft-delete charts whose only remaining dashboard is the one being deleted. List filteredChartRecordsToBeDeleted = new ArrayList<>(); @@ -277,13 +289,12 @@ protected void deleteChildren( deleteChildren(filteredChartRecordsToBeDeleted, hardDelete, updatedBy); } - // Override restore behavior to handle charts through HAS relation. + // Restore chart links (HAS relation). The CONTAINS subtree is now restored by the bulk + // path in EntityRepository.bulkRestoreSubtree; chart handling is a per-dashboard concern + // and lives in the per-entity extension hook. @Transaction @Override - protected void restoreChildren(UUID dashboardId, String updatedBy) { - super.restoreChildren(dashboardId, updatedBy); - - // Load all charts linked to this dashboard + protected void restoreAdditionalChildren(UUID dashboardId, String updatedBy) { List chartRecords = daoCollection .relationshipDAO() @@ -292,7 +303,6 @@ protected void restoreChildren(UUID dashboardId, String updatedBy) { return; } - // Batch-load dashboard relationships for these charts List dashboardRelationships = daoCollection .relationshipDAO() @@ -315,11 +325,9 @@ protected void restoreChildren(UUID dashboardId, String updatedBy) { Include.DELETED) .stream() .map(Dashboard::getId) - .filter(id -> !id.equals(dashboardId)) // (excluding the current dashboard + .filter(id -> !id.equals(dashboardId)) .collect(Collectors.toSet()); - // For restore: get charts whose linked dashboards (excluding the current dashboard) - // are all non‑deleted. List filteredChartRecordsToBeRestored = new ArrayList<>(); @@ -341,9 +349,25 @@ protected void restoreChildren(UUID dashboardId, String updatedBy) { } } + // Per-chart restore preserves the full chart restoreEntity flow (setFieldsInternal, + // setInheritedFields, lifecycle hooks, ES restore-from-search). Charts are typically + // few per dashboard, so the loop isn't a hot path; the bulkRestoreSubtree shortcut + // skipped chart-specific setup that the test in DashboardResourceIT relies on. for (CollectionDAO.EntityRelationshipRecord record : filteredChartRecordsToBeRestored) { LOG.info("Recursively restoring {} {}", record.getType(), record.getId()); - Entity.restoreEntity(updatedBy, record.getType(), record.getId()); + try { + Entity.restoreEntity(updatedBy, record.getType(), record.getId()); + } catch (RuntimeException e) { + // Surface the underlying cause — Entity.restoreEntity has no try/catch wrapper of + // its own and silently aborts the whole dashboard restore if a single chart fails. + LOG.error( + "[ChartRestoreCascade] Failed to restore chart {} for dashboard {}: {}", + record.getId(), + dashboardId, + e.getMessage(), + e); + throw e; + } } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityDAO.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityDAO.java index 6de4d7cd3c86..287c3188f579 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityDAO.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityDAO.java @@ -15,7 +15,8 @@ import static org.openmetadata.service.exception.CatalogExceptionMessage.entityNotFound; import static org.openmetadata.service.jdbi3.ListFilter.escape; -import static org.openmetadata.service.jdbi3.ListFilter.escapeApostrophe; +import static org.openmetadata.service.jdbi3.ListFilter.escapeBackslashAndApostrophe; +import static org.openmetadata.service.jdbi3.ListFilter.escapeForMySqlRegexReplacement; import static org.openmetadata.service.jdbi3.locator.ConnectionType.MYSQL; import static org.openmetadata.service.jdbi3.locator.ConnectionType.POSTGRES; @@ -56,6 +57,22 @@ public interface EntityDAO { org.slf4j.Logger LOG = org.slf4j.LoggerFactory.getLogger(EntityDAO.class); + /** + * Maximum number of values expanded into a single SQL IN-list. JDBI's {@code @BindList} + * produces one bind parameter per element. OpenMetadata supports MySQL and PostgreSQL — + * PostgreSQL's protocol caps each statement at 65535 bind parameters + * (the {@code int2}-size {@code numParams} field), and MySQL's {@code max_allowed_packet} + * caps total statement size. 30k UUID/hash strings stays comfortably under both: each + * UUID is ~36 chars, so an IN-list of this size is ~1MB on the wire (well below the 64MB + * MySQL default) and still leaves headroom for Postgres's parameter ceiling. Callers that + * may exceed this size must chunk their input lists; helpers in this interface + * ({@link #findEntitiesByIds}, {@link #findEntityByNames}, {@link #findReferencesByFqns}, + * {@link #deleteByIds}) already do. (SQL Server isn't a supported connection type here — + * its ~2100 sp_executesql cap would require a separate, much smaller constant if it ever + * is.) + */ + int MAX_IN_LIST_CHUNK_SIZE = 30_000; + /** Methods that need to be overridden by interfaces extending this */ String getTableName(); @@ -272,7 +289,7 @@ default List findReferencesByFqns(List entityFQNs, Incl } List nameHashes = entityFQNs.stream().distinct().map(FullyQualifiedName::buildHash).toList(); - int maxChunkSize = 30000; + int maxChunkSize = MAX_IN_LIST_CHUNK_SIZE; if (nameHashes.size() <= maxChunkSize) { return findReferenceRows(nameHashes, include).stream() .map(row -> row.toEntityReference(Entity.getEntityTypeFromClass(getEntityClass()))) @@ -325,6 +342,16 @@ default void updateFqn(String oldPrefix, String newPrefix) { if (!getNameHashColumn().equals("fqnHash")) { return; } + // The regex replacement argument to MySQL's REGEXP_REPLACE has its own escape layer + // on top of the SQL string-literal layer — `\1`/`\2` are backreferences, `\\` is a + // literal backslash. Using escapeBackslashAndApostrophe here would only escape for the + // SQL layer, leaving a stray backslash in newPrefix to be interpreted by the regex + // engine. escapeForMySqlRegexReplacement applies both layers (regex-replacement first, + // then SQL string-literal) so an input backslash round-trips to a single literal + // backslash in the replacement output. The source pattern goes through escape() which + // already covers the SQL + LIKE-underscore layers — the regex-pattern layer is + // tolerated here because OpenMetadata's name validation forbids the regex metas that + // would matter (\ . * ? + ^ $ ( ) [ ] { } |). String mySqlUpdate = String.format( "UPDATE %s SET json = " @@ -333,11 +360,16 @@ default void updateFqn(String oldPrefix, String newPrefix) { + "WHERE fqnHash LIKE '%s.%%'", getTableName(), escape(oldPrefix), - escapeApostrophe(newPrefix), + escapeForMySqlRegexReplacement(newPrefix), FullyQualifiedName.buildHash(oldPrefix), FullyQualifiedName.buildHash(newPrefix), FullyQualifiedName.buildHash(oldPrefix)); + // Postgres path embeds the prefixes inside a double-quoted JSON pattern, so escape + // backslashes and apostrophes first (so a literal "\\" or "''" isn't reparsed by the + // SQL string-literal layer), then escape double-quotes so the JSON-pattern delimiter + // can't be broken out of. Apostrophe escaping is still required because the JSON + // pattern itself sits inside a single-quoted SQL string literal. String postgresUpdate = String.format( "UPDATE %s SET json = " @@ -346,8 +378,8 @@ default void updateFqn(String oldPrefix, String newPrefix) { + ", fqnHash = REPLACE(fqnHash, '%s.', '%s.') " + "WHERE fqnHash LIKE '%s.%%'", getTableName(), - ReindexingUtil.escapeDoubleQuotes(escapeApostrophe(oldPrefix)), - ReindexingUtil.escapeDoubleQuotes(escapeApostrophe(newPrefix)), + ReindexingUtil.escapeDoubleQuotes(escapeBackslashAndApostrophe(oldPrefix)), + ReindexingUtil.escapeDoubleQuotes(escapeBackslashAndApostrophe(newPrefix)), FullyQualifiedName.buildHash(oldPrefix), FullyQualifiedName.buildHash(newPrefix), FullyQualifiedName.buildHash(oldPrefix)); @@ -608,6 +640,26 @@ boolean existsByName( @SqlUpdate("DELETE FROM

WHERE id = :id") int delete(@Define("table") String table, @BindUUID("id") UUID id); + @SqlUpdate("DELETE FROM
WHERE id IN ()") + int deleteByIds(@Define("table") String table, @BindList("ids") List ids); + + default int deleteByIds(List ids) { + if (ids == null || ids.isEmpty()) { + return 0; + } + List stringIds = ids.stream().map(UUID::toString).toList(); + int maxChunkSize = MAX_IN_LIST_CHUNK_SIZE; + if (stringIds.size() <= maxChunkSize) { + return deleteByIds(getTableName(), stringIds); + } + int deleted = 0; + for (int i = 0; i < stringIds.size(); i += maxChunkSize) { + List chunk = stringIds.subList(i, Math.min(i + maxChunkSize, stringIds.size())); + deleted += deleteByIds(getTableName(), chunk); + } + return deleted; + } + @ConnectionAwareSqlUpdate(value = "ANALYZE TABLE
", connectionType = MYSQL) @ConnectionAwareSqlUpdate(value = "ANALYZE
", connectionType = POSTGRES) void analyze(@Define("table") String table); @@ -694,7 +746,7 @@ default List findEntitiesByIds(List ids, Include include) { } List distinctIds = ids.stream().map(UUID::toString).distinct().toList(); - int maxChunkSize = 30000; + int maxChunkSize = MAX_IN_LIST_CHUNK_SIZE; if (distinctIds.size() <= maxChunkSize) { return findByIds(getTableName(), distinctIds, getCondition(include)).stream() @@ -739,7 +791,7 @@ default List findEntityByNames(List entityFQNs, Include include) { } List names = entityFQNs.stream().distinct().map(FullyQualifiedName::buildHash).toList(); - int maxChunkSize = 30000; + int maxChunkSize = MAX_IN_LIST_CHUNK_SIZE; if (names.size() <= maxChunkSize) { return findByNames(getTableName(), getNameHashColumn(), names, getCondition(include)).stream() diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java index 64185467245c..3b18cc25d9f9 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/EntityRepository.java @@ -145,6 +145,7 @@ import java.util.concurrent.TimeUnit; import java.util.function.BiConsumer; import java.util.function.BiPredicate; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Supplier; import java.util.stream.Collectors; @@ -4260,7 +4261,16 @@ private DeleteResponse delete( EntityUpdater updater = getUpdater(original, updated, Operation.SOFT_DELETE, null); updater.update(); changeType = ENTITY_SOFT_DELETED; + // Run the same hook the bulk path runs — keeps direct-entity soft delete in sync + // with bulkSoftDeleteSubtree for repos that link non-CONTAINS entities (e.g., + // dashboard charts). + softDeleteAdditionalChildren(original.getId(), deletedBy); } else { + // Run hook BEFORE cleanup(): cleanup() deletes this entity's relationship rows + // (including HAS), and subclass hooks like DashboardRepository.cascadeChartCleanup + // need to walk HAS to discover linked entities. Mirrors bulkHardDeleteSubtree + // ordering for direct-entity hard delete. + hardDeleteAdditionalChildren(original.getId(), deletedBy); cleanup(updated); changeType = ENTITY_DELETED; } @@ -4336,153 +4346,25 @@ protected void deleteChildren(UUID id, boolean recursive, boolean hardDelete, St @Transaction protected void deleteChildren( List children, boolean hardDelete, String updatedBy) { - // Use batch deletion only for hard deletes with large numbers of children - // For soft deletes, we must maintain the correct order for restoration to work properly - if (hardDelete && children.size() > 100) { - LOG.info("Using batch deletion for {} children entities", children.size()); - batchDeleteChildren(children, hardDelete, updatedBy); - } else { - // For soft deletes or small numbers, use original sequential deletion - // This ensures proper parent-child relationships are maintained for restoration - for (EntityRelationshipRecord entityRelationshipRecord : children) { - LOG.info( - "Recursively {} deleting {} {}", - hardDelete ? "hard" : "soft", - entityRelationshipRecord.getType(), - entityRelationshipRecord.getId()); - Entity.deleteEntity( - updatedBy, - entityRelationshipRecord.getType(), - entityRelationshipRecord.getId(), - true, - hardDelete); - } + if (children.isEmpty()) { + return; } - } - - /** - * Batch deletion of children entities for improved performance - */ - @Transaction - protected void batchDeleteChildren( - List children, boolean hardDelete, String updatedBy) { - - // Group entities by type for batch processing - Map> entitiesByType = + // Both soft-delete and hard-delete dispatch to the per-type bulk path. One batched DB + // write + one batched change-event insert per type, regardless of descendant count. + // For hard delete, bulkHardDeleteSubtree replaces the legacy per-entity cleanup loop + // that opened an independent JDBI transaction per descendant. + Map> idsByType = children.stream() .collect( Collectors.groupingBy( EntityRelationshipRecord::getType, Collectors.mapping(EntityRelationshipRecord::getId, Collectors.toList()))); - - LOG.info("Batch deleting {} entities across {} types", children.size(), entitiesByType.size()); - - // Process deletion in levels to handle cascading properly - for (Map.Entry> entry : entitiesByType.entrySet()) { - String childEntityType = entry.getKey(); - List entityIds = entry.getValue(); - - LOG.info("Batch processing {} entities of type {}", entityIds.size(), childEntityType); - - // Process in smaller batches to avoid overwhelming the system - int batchSize = 50; - for (int i = 0; i < entityIds.size(); i += batchSize) { - List batch = entityIds.subList(i, Math.min(i + batchSize, entityIds.size())); - processDeletionBatch(batch, childEntityType, hardDelete, updatedBy); - } - } - } - - /** - * Process a batch of entities for hard deletion. Entered only via - * {@link #batchDeleteChildren}, which only fires for {@code hardDelete=true} and - * {@code children.size() > 100}; the soft-delete and small-batch paths stay on the - * sequential {@link Entity#deleteEntity} flow. - * - *

Each child is removed via {@link #cleanup}, which deletes the entity row, all - * {@code (id, *)} and {@code (*, id)} entity_relationship rows, extensions, tag usage, - * threads, and caches as one atomic unit per child (cleanup opens its own JDBI - * transaction via {@code Entity.getJdbi().inTransaction(...)}). The previous - * implementation pre-deleted relationships in two batched queries before this loop and - * swallowed exceptions in the per-child cleanup, which is exactly what produced the - * orphan pattern: a failed cleanup left an entity row alive after its relationships - * had been wiped, surfacing in {@code /containers?root=true} and breaking - * {@code /children} traversal. - * - *

Failure semantics: per-child atomicity, not whole-batch atomicity. If - * cleanup for child k throws, children {@code 0..k-1} have already committed - * via their own transactions and cannot be rolled back by the {@code @Transaction} on - * this method (cleanup's inner transaction is independent). The exception propagates - * so the loop stops; children {@code k..N-1} keep both their rows and their parent - * CONTAINS relationships intact. The retry path is to reissue the recursive delete on - * the parent — remaining children re-enter this loop. Crucially, no orphan-without- - * relationships row can result from an exception in this method, which is the bug - * this change exists to fix; achieving true all-or-nothing rollback across the batch - * would require sharing one JDBI handle across every cleanup call, a wider refactor - * deliberately scoped out of this fix. - */ - @Transaction - private void processDeletionBatch( - List entityIds, String entityType, boolean hardDelete, String updatedBy) { - - LOG.debug("Processing batch of {} {} entities", entityIds.size(), entityType); - - // First, collect all grandchildren that need to be deleted in a SINGLE batch query - List stringIds = entityIds.stream().map(UUID::toString).collect(Collectors.toList()); - List grandchildRecords = - daoCollection - .relationshipDAO() - .findToBatchWithRelations( - stringIds, - entityType, - List.of(Relationship.CONTAINS.ordinal(), Relationship.PARENT_OF.ordinal())); - - // Convert to EntityRelationshipRecord format - List allGrandchildren = - grandchildRecords.stream() - .map( - rec -> - new EntityRelationshipRecord( - UUID.fromString(rec.getToId()), rec.getToEntity(), rec.getJson())) - .collect(Collectors.toList()); - - // Recursively delete grandchildren first - if (!allGrandchildren.isEmpty()) { - LOG.info("Found {} grandchildren to delete first", allGrandchildren.size()); - deleteChildren(allGrandchildren, hardDelete, updatedBy); - } - - // cleanup() per entity is the source of truth: it removes the row, its relationships - // (deleteAllFrom + deleteAllTo on (id, *) and (*, id)), extensions, tag usage, feed - // threads, and caches within ONE JDBI transaction owned by cleanup itself. The - // previous pre-batch-delete of relationships made the row-and-relationship pairing - // non-atomic across the loop, which is why a swallowed mid-loop exception produced - // the orphan-without-relationships pattern. Letting cleanup own both halves and - // letting exceptions propagate stops the loop early on failure; see the failure - // semantics note in the Javadoc above for what "stops the loop" actually guarantees. - @SuppressWarnings("rawtypes") - EntityRepository repository = Entity.getEntityRepository(entityType); - for (UUID entityId : entityIds) { - try { - EntityInterface entity = repository.find(entityId, Include.ALL); - repository.cleanup(entity); - } catch (RuntimeException e) { - LOG.error( - "Failed to delete {} '{}' during recursive batch delete: {}", - entityType, - entityId, - e.getMessage(), - e); - // Wrap with entity context before re-throwing so the operator can identify - // the row that blocked a large recursive delete. The exception still - // propagates — the loop still stops, the failure-semantics contract in the - // Javadoc still holds — we just trade an opaque stack trace for one that - // names the offending child. - throw new RuntimeException( - String.format( - "Failed to delete %s '%s' during recursive batch delete: %s", - entityType, entityId, e.getMessage()), - e); + for (var entry : idsByType.entrySet()) { + EntityRepository repo = Entity.getEntityRepository(entry.getKey()); + if (hardDelete) { + repo.bulkHardDeleteSubtree(entry.getValue(), updatedBy); + } else { + repo.bulkSoftDeleteSubtree(entry.getValue(), updatedBy); } } } @@ -5683,15 +5565,31 @@ public final URI getHref(UriInfo uriInfo, UUID id) { @Transaction public final PutResponse restoreEntity(String updatedBy, UUID id) { + // Confirm the entity exists at all (in any state). If the row is truly gone + // (e.g., hard-deleted), propagate EntityNotFoundException so the caller surfaces + // a clean 404 instead of running children / hooks against a non-existent id and + // potentially surfacing a 500 from a hook side-effect. + find(id, ALL); + // If an entity being restored contains other **deleted** children entities, restore them restoreChildren(id, updatedBy); // Finally set entity deleted flag to false LOG.info("Restoring the {} {}", entityType, id); + PutResponse response = null; try { T original = find(id, DELETED); - setFieldsInternal(original, putFields); + // Populate fields with Include.ALL so HAS-style children that were soft-deleted as + // part of this entity's cascade remain in the loaded child lists (e.g., + // dashboard.charts). If we used the default NON_DELETED filter, those lists would + // come back empty, and the PUT updater's diff would see "no children" on both + // sides and call deleteFrom(...) to wipe every HAS relationship row before the + // additional-children hook ever runs to restore them. Charts attached to the + // dashboard being restored would then have nothing to walk back from, and the + // restore cascade would silently no-op (DashboardResourceIT#test_deleteDashboard_ + // chartBelongsToSingleDashboard_chartIsDeletedThenRestored guards against this). + setFieldsInternal(original, putFields, ALL); setInheritedFields(original, putFields); T updated = JsonUtils.readValue(JsonUtils.pojoToJson(original), entityClass); updated.setUpdatedBy(updatedBy); @@ -5700,24 +5598,578 @@ public final PutResponse restoreEntity(String updatedBy, UUID id) { updater.update(); // Restore moves the row from deleted=true to deleted=false, changing the listing total. ListCountCache.invalidate(entityType); - return new PutResponse<>(Status.OK, updated, ENTITY_RESTORED); + response = new PutResponse<>(Status.OK, updated, ENTITY_RESTORED); } catch (EntityNotFoundException e) { - LOG.info("Entity is not in deleted state {} {}", entityType, id); - return null; + // Entity exists (verified above) but is not in DELETED state — already restored. + LOG.info("Entity already restored or not in deleted state {} {}", entityType, id); } + // Run the per-entity hook because the entity exists (the find(ALL) guard ensures + // that). A re-entered cascade where this level is already restored must still + // reconcile HAS-related children (e.g., dashboard charts) of nested descendants. + restoreAdditionalChildren(id, updatedBy); + return response; } @Transaction protected void restoreChildren(UUID id, String updatedBy) { - // Restore deleted children entities + // Walk CONTAINS + PARENT_OF so the restore cascade is symmetric with deleteChildren + // and the bulk subtree walkers — Team → Team, KnowledgePage → KnowledgePage, + // Classification → Tag etc. express their hierarchy via PARENT_OF, and a CONTAINS-only + // probe would skip them on restore even though delete already cascades through them. List records = - daoCollection.relationshipDAO().findTo(id, entityType, Relationship.CONTAINS.ordinal()); - if (!records.isEmpty()) { - // Recursively restore all contained entities - for (CollectionDAO.EntityRelationshipRecord record : records) { - LOG.info("Recursively restoring {} {}", record.getType(), record.getId()); - Entity.restoreEntity(updatedBy, record.getType(), record.getId()); + daoCollection + .relationshipDAO() + .findTo( + id, + entityType, + List.of(Relationship.CONTAINS.ordinal(), Relationship.PARENT_OF.ordinal())); + if (records.isEmpty()) { + return; + } + Map> idsByType = new HashMap<>(); + for (CollectionDAO.EntityRelationshipRecord record : records) { + idsByType.computeIfAbsent(record.getType(), k -> new ArrayList<>()).add(record.getId()); + } + for (var entry : idsByType.entrySet()) { + EntityRepository repo = Entity.getEntityRepository(entry.getKey()); + repo.bulkRestoreSubtree(entry.getValue(), updatedBy); + } + } + + /** + * Bulk-restore a set of soft-deleted entities of this repository's type along with their entire + * subtree of CONTAINS-related descendants. Replaces the per-entity recursive path that was + * O(descendants) HTTP-request-bound work with a per-level batched walk that uses the existing + * deferred-store bulk update infrastructure. + * + *

For a database with N descendants, the previous implementation issued ~N find calls, + * ~N updates and ~N search index writes, all serialized inside one HTTP request. This path + * does one batched DB load, one batched DB write and one batched change-event insert per + * level, and relies on {@link #restoreFromSearch(EntityInterface)} at the top-level to + * cascade the deleted flag flip across child indexes in a single ES update_by_query. + * + *

Subclasses that link non-CONTAINS related entities (e.g., charts attached to dashboards + * via HAS) should implement the {@link #restoreAdditionalChildren(UUID, String)} hook — + * the CONTAINS subtree is restored by the bulk path itself, so per-entity overrides of + * {@code restoreChildren} are no longer invoked from inside the bulk walk. + * + *

Operational ceiling: the entire walk runs inside a single JDBI + * {@code @Transaction}, which holds one connection from the pool for the duration. The + * async restore endpoint ({@code ?async=true}) moves the work onto a virtual thread but + * keeps the same single-transaction shape — it just lets the client get a 202 back. + * Back-pressure under load comes from the JDBI connection pool itself: virtual threads + * are cheap, so under saturation tasks queue on connection acquisition (with the pool's + * own timeout) rather than at the executor. Chunked-transaction support is tracked as a + * follow-up if this becomes a real bottleneck. + */ + @Transaction + public final void bulkRestoreSubtree(List ids, String updatedBy) { + if (ids == null || ids.isEmpty()) { + return; + } + // Load with ALL — we still need to walk children when the parents at this level are + // already restored (or never deleted), in case deeper descendants are deleted and + // must be flipped. Matches the previous recursive path that always called + // restoreChildren before checking the parent's deleted state. + List entities = loadForBulk(ids, ALL, "bulkRestoreLoad"); + if (entities.isEmpty()) { + return; + } + dispatchToContainedChildren( + entities, + "bulkRestoreFindChildren", + (childRepo, childIds) -> childRepo.bulkRestoreSubtree(childIds, updatedBy)); + + List deletedEntities = + entities.stream().filter(e -> Boolean.TRUE.equals(e.getDeleted())).toList(); + if (!deletedEntities.isEmpty()) { + // Hydrate relationship fields with Include.ALL before the PUT updater diff runs. + // loadForBulk returned only the storage JSON, so HAS-style children + // (e.g., dashboard.charts, dashboard.dataModels) are null on the parsed entity. + // The PUT updater's compareAndUpdate("charts", ...) fires unconditionally and the + // update(...) lambda does deleteFrom(... HAS ...) followed by re-adding from + // updated.getCharts() — if updated.getCharts() is null/empty, every HAS row is + // wiped before the restoreAdditionalChildren hook ever runs to restore them. + // Using Include.ALL ensures the cascade-deleted charts/dataModels are visible to + // both sides of the diff so the relationships round-trip cleanly. Matches the + // single-entity restoreEntity contract (see the comment at the find/setFields call + // earlier in this file). + hydrateRelationsForBulkUpdater(deletedEntities); + List updaters = + buildBulkUpdaters(deletedEntities, updatedBy, Operation.PUT, "bulkRestoreUpdaters", null); + List changed = filterChanged(updaters); + if (!changed.isEmpty()) { + persistBulkUpdaters(changed, ENTITY_RESTORED, updatedBy, "bulkRestore"); + ListCountCache.invalidate(entityType); + } + } + // Always run per-entity hooks even when nothing at THIS level needed flipping — + // a re-entered cascade may still have HAS-related children attached to nested + // descendants that require reconciliation. + runRestoreAdditionalChildren(entities, updatedBy); + } + + private void runRestoreAdditionalChildren(List entities, String updatedBy) { + for (T entity : entities) { + restoreAdditionalChildren(entity.getId(), updatedBy); + } + } + + /** + * Default relation set walked when descending into a parent's subtree. CONTAINS covers the + * service → DB → schema → table chain and most other parent → child hierarchies; PARENT_OF + * covers recursive shapes like Glossary → GlossaryTerm, Team → Team, Classification → Tag, + * Domain → DataProduct. Walking both keeps every entity type's hierarchy in scope without + * subclass-specific overrides. + */ + private static final List SUBTREE_RELATIONS = + List.of(Relationship.CONTAINS.ordinal(), Relationship.PARENT_OF.ordinal()); + + /** + * Find all subtree children (CONTAINS + PARENT_OF) for every entity in {@code parents} with one + * batched query, then apply {@code dispatcher} to each (childRepo, childIds) group. Replaces the + * per-parent {@code findTo} round-trip that used to fire once per descendant — for a 12k-table + * database that's 12k DB hits collapsed into one per tree level. Shared between bulk restore, + * bulk soft-delete and bulk hard-delete; the only thing that varies is the terminal call on the + * child repo. + */ + private void dispatchToContainedChildren( + List parents, String phaseName, BiConsumer, List> dispatcher) { + List parentIds = new ArrayList<>(parents.size()); + for (T parent : parents) { + parentIds.add(parent.getId().toString()); + } + List relationships; + try (var ignored = phase(phaseName)) { + relationships = + daoCollection.relationshipDAO().findToBatchAllTypes(parentIds, SUBTREE_RELATIONS, ALL); + } + if (relationships.isEmpty()) { + return; + } + Map> idsByChildType = new HashMap<>(); + for (var rel : relationships) { + if (!entityType.equals(rel.getFromEntity())) { + continue; + } + idsByChildType + .computeIfAbsent(rel.getToEntity(), k -> new ArrayList<>()) + .add(UUID.fromString(rel.getToId())); + } + for (var entry : idsByChildType.entrySet()) { + EntityRepository repo = Entity.getEntityRepository(entry.getKey()); + dispatcher.accept(repo, entry.getValue()); + } + } + + /** + * Hook called once per restored entity for repositories that have non-CONTAINS related + * entities that need to be restored alongside the parent. Default: no-op. + */ + protected void restoreAdditionalChildren(UUID id, String updatedBy) { + // No-op. Override in subclasses for HAS-style related-entity restore. + } + + /** + * Bulk soft-delete the given entities of this repository's type along with their CONTAINS + * subtree. Symmetric to {@link #bulkRestoreSubtree(List, String)}: replaces the per-entity + * recursive {@code Entity.deleteEntity} loop in + * {@link #deleteChildren(List, boolean, String)} with a per-level batched walk that uses + * the deferred-store bulk update infrastructure. + * + *

Per-level shape: one batched {@code findToBatchAllTypes}, one batched DB load (NON + * deleted only — already-deleted entities are skipped, mirroring the per-entity guard), + * one batched {@code updateMany} that flips {@code deleted = true}, one batched version + * history insert, one batched change-event insert, one batched cache invalidation. + * Per-descendant ES writes are skipped — the top-level + * {@link #deleteFromSearch(EntityInterface, boolean)} cascade flips the deleted flag on + * descendant ES indexes in a single update_by_query. + * + *

Entity types where {@code supportsSoftDelete} is false fall back to the per-entity + * hard-delete path (matches the existing per-entity {@code delete()} fallback). Subclasses + * with non-CONTAINS linked entities should override + * {@link #softDeleteAdditionalChildren(UUID, String)}. + * + *

Operational ceiling: see {@link #bulkRestoreSubtree(List, String)} — the same + * single-{@code @Transaction} shape applies on the delete side. Chunked-transaction + * support is tracked as a follow-up. + */ + @Transaction + public final void bulkSoftDeleteSubtree(List ids, String updatedBy) { + if (ids == null || ids.isEmpty()) { + return; + } + if (!supportsSoftDelete) { + hardDeleteAtLevelOnly(ids, updatedBy); + return; + } + List allEntities = loadForBulk(ids, ALL, "bulkSoftDeleteLoad"); + if (allEntities.isEmpty()) { + return; + } + List entities = + allEntities.stream().filter(e -> !Boolean.TRUE.equals(e.getDeleted())).toList(); + for (T entity : entities) { + checkSystemEntityDeletion(entity); + preDelete(entity, updatedBy); + } + dispatchToContainedChildren( + allEntities, + "bulkSoftDeleteFindChildren", + (childRepo, childIds) -> childRepo.bulkSoftDeleteSubtree(childIds, updatedBy)); + applyBulkSoftDelete(entities, updatedBy); + // Always run per-entity hooks even when nothing at THIS level needed flipping — + // descendants restored independently before the cascade still need to be re-deleted + // by the per-entity hook. + runSoftDeleteAdditionalChildren(allEntities, updatedBy); + } + + // This type can't be soft-deleted, so each entity at this level must be hard + // deleted instead. Pass hardDelete=false through to the per-entity delete so + // descendant levels that *do* support soft delete remain soft-deleted — the + // per-entity flow handles the asymmetry by inspecting each level's own + // supportsSoftDelete flag. + private void hardDeleteAtLevelOnly(List ids, String updatedBy) { + for (UUID id : ids) { + Entity.deleteEntity(updatedBy, entityType, id, true, false); + } + } + + private void applyBulkSoftDelete(List entities, String updatedBy) { + if (entities.isEmpty()) { + return; + } + // Same reason as hydrateRelationsForBulkUpdater — buildBulkUpdaters uses bare JSON, and a + // PUT-style updater (e.g. DashboardUpdater.entitySpecificUpdate) calls + // deleteFrom(... HAS ...) then re-adds from updated.getCharts(). Without hydration + // both lists are empty and the soft-delete wipes the HAS rows that softDeleteAdditional- + // Children later needs to walk. Include.ALL handles both shapes: charts that are still + // live (parent soft-deleted in isolation) and charts already cascade-soft-deleted + // (parent soft-deleted as part of a wider sweep). + hydrateRelationsForBulkUpdater(entities); + List updaters = + buildBulkUpdaters( + entities, + updatedBy, + Operation.SOFT_DELETE, + "bulkSoftDeleteUpdaters", + e -> e.setDeleted(true)); + List changed = filterChanged(updaters); + if (!changed.isEmpty()) { + persistBulkUpdaters(changed, ENTITY_SOFT_DELETED, updatedBy, "bulkSoftDelete"); + ListCountCache.invalidate(entityType); + } + } + + private void runSoftDeleteAdditionalChildren(List entities, String updatedBy) { + for (T entity : entities) { + softDeleteAdditionalChildren(entity.getId(), updatedBy); + } + } + + /** + * Hook called once per soft-deleted entity for repositories that have non-CONTAINS related + * entities that need to be soft-deleted alongside the parent (e.g., charts attached to + * dashboards via HAS). Default: no-op. + */ + protected void softDeleteAdditionalChildren(UUID id, String updatedBy) { + // No-op. Override in subclasses for HAS-style related-entity soft delete. + } + + /** + * Bulk hard-delete the given entities of this repository's type along with their entire + * CONTAINS + PARENT_OF subtree. Replaces the legacy per-entity {@link #cleanup} loop driven by + * {@code processDeletionBatch} / {@code batchDeleteChildren} — that path opened an independent + * JDBI transaction per descendant and fired ~10 SQL statements per entity, so a 12k-table + * database needed ~120,000 round-trips and produced the hours-long deletes reported by users. + * + *

Per-level shape: one batched {@code findToBatchAllTypes} that walks both CONTAINS + * (service → DB → schema → table) and PARENT_OF (Glossary → GlossaryTerm, Team → Team, recursive + * Container) so every entity hierarchy is in scope without per-subclass overrides; one batched + * DB load; recursive descent into each child type; one + * {@link CollectionDAO.EntityRelationshipDAO#batchDeleteRelationships} per type to wipe both + * {@code (id, *)} and {@code (*, id)} entity_relationship rows in a single statement; one + * batched extension delete; one batched entity row delete; per-entity loops for tag_usage / + * usage / field_relationship / feed threads (those tables key on FQN strings rather than ids + * so they can't share a single IN-list query, but they stay inside the same {@code @Transaction} + * which removes the per-entity transaction overhead that dominated the old path). + * + *

Subclasses with non-CONTAINS related entities (e.g., dashboard charts attached via HAS) + * should override {@link #hardDeleteAdditionalChildren(UUID, String)}. Subclasses that need true + * batched external cleanup (Airflow DAGs, S3, secrets stores) can override + * {@link #bulkEntitySpecificCleanup(List)}; the default loops the per-entity hook. + * + *

Failure semantics: the entire bulk hard-delete runs in a single + * {@code @Transaction}, so a mid-walk failure rolls back every row + relationship deletion. + * This is stronger than the previous {@code processDeletionBatch} contract, which only + * guaranteed per-child atomicity and could leave the operator with a partially-deleted subtree + * after a failure. See also {@link #bulkRestoreSubtree(List, String)} for the same operational + * ceiling note around single-connection holding for the duration of the walk. + */ + @Transaction + public final void bulkHardDeleteSubtree(List ids, String updatedBy) { + if (ids == null || ids.isEmpty()) { + return; + } + List entities = loadForBulk(ids, ALL, "bulkHardDeleteLoad"); + if (entities.isEmpty()) { + return; + } + // Populate relation fields up front so the same subclass hooks the legacy + // Entity.deleteEntity path called against a fully-loaded entity (e.g., + // TestCaseRepository.updateTestSuite reading testCase.getTestSuite()) see the + // expected shape. bulkCleanupReferences wipes these relationship rows later, so + // hooks running after that point must remain null-safe. + populateRelationFields(entities); + for (T entity : entities) { + checkSystemEntityDeletion(entity); + preDelete(entity, updatedBy); + } + dispatchToContainedChildren( + entities, + "bulkHardDeleteFindChildren", + (childRepo, childIds) -> childRepo.bulkHardDeleteSubtree(childIds, updatedBy)); + bulkEntitySpecificCleanup(entities); + // Run BEFORE bulkCleanupReferences: hooks like DashboardRepository.cascadeChartCleanup + // walk HAS relationships to discover linked entities, and bulkCleanupReferences wipes + // those relationship rows. + runHardDeleteAdditionalChildren(entities, updatedBy); + bulkCleanupReferences(entities); + bulkDeleteEntityRows(entities); + bulkInvalidate(entities); + for (T entity : entities) { + postDelete(entity, true); + // Fire deleteFromSearch per-entity so cascade-deleted descendants are removed from + // Elasticsearch. The legacy per-entity Entity.deleteEntity path invoked this via + // delete()'s top-level dispatch — this bulk replacement is the only path that walks + // cascaded children now, so a missing call leaves stale ES docs that surface as + // duplicate results (e.g. Playwright Domains.spec.ts:533 found two "PW_DataProduct_ + // Sales" rows after a recursive Domain hard-delete because the DB row was gone but + // the search-index doc lingered). + deleteFromSearch(entity, true); + } + } + + private void populateRelationFields(List entities) { + try { + setFieldsInBulk(putFields, entities); + } catch (Exception e) { + LOG.debug( + "Bulk field population failed during bulk hard delete for {}, falling back per-entity: {}", + entityType, + e.getMessage()); + for (T entity : entities) { + try { + setFieldsInternal(entity, putFields); + } catch (Exception ignored) { + // postDelete subclass overrides must remain null-safe for cascade-deleted parents. + } + } + } + } + + /** + * Per-entity hydration with {@link Include#ALL} for the bulk restore path. The bulk + * {@link #setFieldsInBulk} variant hard-codes {@code NON_DELETED} when batch-fetching + * relationship references (see {@code DashboardRepository.batchFetchCharts}), so a + * cascade-deleted chart wouldn't show up in {@code dashboard.charts} — exactly the + * scenario where we need it to. Falling back to per-entity {@link #setFieldsInternal} + * routes through the subclass's {@code setFields(entity, fields, relationIncludes)} which + * honours the include passed in. Restore batches are typically small (single subtree + * level), so the extra DB round-trips are acceptable for the correctness this buys. + */ + private void hydrateRelationsForBulkUpdater(List entities) { + for (T entity : entities) { + try { + setFieldsInternal(entity, putFields, ALL); + } catch (Exception ex) { + // Best-effort: if hydration fails on a single entity the PUT updater may wipe its + // HAS rows. restoreAdditionalChildren will still attempt to put them back, but log + // so operators can correlate any missing-relationship reports with hydration noise + // rather than digging through change-event history. + LOG.warn( + "Hydration failed for {} {}; HAS rows may be wiped before restore hook runs", + entityType, + entity.getId(), + ex); + } + } + } + + private void bulkCleanupReferences(List entities) { + List entityIds = new ArrayList<>(entities.size()); + List entityIdStrings = new ArrayList<>(entities.size()); + for (T entity : entities) { + entityIds.add(entity.getId()); + entityIdStrings.add(entity.getId().toString()); + } + try (var ignored = phase("bulkHardDeleteRelationships")) { + daoCollection.relationshipDAO().batchDeleteRelationships(entityIds, entityType); + } + try (var ignored = phase("bulkHardDeleteExtensions")) { + daoCollection.entityExtensionDAO().deleteAllBatch(entityIdStrings); + } + try (var ignored = phase("bulkHardDeleteFqnDependents")) { + for (T entity : entities) { + String fqn = entity.getFullyQualifiedName(); + daoCollection.fieldRelationshipDAO().deleteAllByPrefix(fqn); + daoCollection.tagUsageDAO().deleteTagLabelsByTargetPrefix(fqn); + daoCollection.tagUsageDAO().deleteTagLabelsByFqn(fqn); + } + } + try (var ignored = phase("bulkHardDeleteUsage")) { + for (T entity : entities) { + daoCollection.usageDAO().delete(entity.getId()); + } + } + try (var ignored = phase("bulkHardDeleteFeedThreads")) { + Entity.getFeedRepository().deleteByAbout(entityIds); + } + } + + private void bulkDeleteEntityRows(List entities) { + try (var ignored = phase("bulkHardDeleteRows")) { + List entityIds = new ArrayList<>(entities.size()); + for (T entity : entities) { + entityIds.add(entity.getId()); + } + dao.deleteByIds(entityIds); + } + } + + private void bulkInvalidate(List entities) { + for (T entity : entities) { + invalidate(entity); + // Mirror cleanup()'s NotFoundCache marker so a concurrent reader that re-populates + // L1/Redis between bulkDeleteEntityRows and the next invalidate doesn't keep + // returning a stale "found" entity. Without this the next get_by_name/find against + // the same id or FQN can still hit the cache and return a deleted entity, which + // breaks fixture teardown (DELETE returns 404 because the row is gone but Redis + // still hands out the entity to the get_by_name probe). + markEntityNotFound(entity); + } + } + + private void runHardDeleteAdditionalChildren(List entities, String updatedBy) { + for (T entity : entities) { + hardDeleteAdditionalChildren(entity.getId(), updatedBy); + } + } + + /** + * Hook called once per hard-deleted entity for repositories that have non-CONTAINS related + * entities that need to be hard-deleted alongside the parent (e.g., charts attached to + * dashboards via HAS). Default: no-op. + */ + protected void hardDeleteAdditionalChildren(UUID id, String updatedBy) { + // No-op. Override in subclasses for HAS-style related-entity hard delete. + } + + /** + * Hook for entity-type-specific cleanup invoked once per bulk-hard-delete batch. Default + * implementation loops {@link #entitySpecificCleanup(EntityInterface)} so subclasses keep + * current behavior. Override for true batching where external resources warrant it (e.g., + * Airflow DAG deregistration, S3 object cleanup, secrets-store purges). + */ + protected void bulkEntitySpecificCleanup(List entities) { + for (T entity : entities) { + entitySpecificCleanup(entity); + } + } + + // ---- Shared phase helpers used by bulkRestoreSubtree / bulkSoftDeleteSubtree ---- + + private List loadForBulk(List ids, Include include, String phaseName) { + try (var ignored = phase(phaseName)) { + return find(ids, include); + } + } + + private List buildBulkUpdaters( + List originals, String updatedBy, Operation op, String phaseName, Consumer mutator) { + long now = System.currentTimeMillis(); + List updaters = new ArrayList<>(originals.size()); + try (var ignored = phase(phaseName)) { + for (T original : originals) { + T updated = JsonUtils.readValue(JsonUtils.pojoToJson(original), entityClass); + updated.setUpdatedBy(updatedBy); + updated.setUpdatedAt(now); + if (mutator != null) { + mutator.accept(updated); + } + EntityUpdater updater = getUpdater(original, updated, op, null); + updater.updateWithDeferredStore(); + updaters.add(updater); + } + } + return updaters; + } + + private List filterChanged(List updaters) { + return updaters.stream().filter(u -> u.isVersionChanged() || u.isEntityChanged()).toList(); + } + + /** + * Apply a batch of {@link EntityUpdater}s already in deferred-store state: write version + * history, persist entity rows, invalidate caches, dispatch the bulk lifecycle event so + * the search index handler updates ES, then emit change events. {@code phasePrefix} is + * used to tag latency phases (e.g. {@code "bulkRestore"} → + * {@code "bulkRestoreVersionHistory"}). + * + *

The lifecycle dispatch is required because the top-level + * {@code restoreFromSearch}/{@code deleteFromSearch} cascade only flips the deleted flag on + * child indexes whose docs join on the parent's id field. HAS-style descendants (e.g., + * charts attached to dashboards) and entity types without a {@code parent.id} field in + * their ES mapping would otherwise drift — DB shows restored / soft-deleted, but ES still + * reflects the previous state. {@code SearchIndexHandler.onEntitiesUpdated} batches the + * writes via {@code updateEntitiesIndex}, so this is still bulk on the ES side. + */ + private void persistBulkUpdaters( + List changed, EventType eventType, String userName, String phasePrefix) { + writeBulkVersionHistory(changed, phasePrefix); + List changedEntities = changed.stream().map(EntityUpdater::getUpdated).toList(); + try (var ignored = phase(phasePrefix + "UpdateMany")) { + updateMany(changedEntities); + } + try (var ignored = phase(phasePrefix + "Invalidate")) { + invalidateMany(changedEntities); + } + try (var ignored = phase(phasePrefix + "LifecycleDispatch")) { + EntityLifecycleEventDispatcher.getInstance().onEntitiesUpdated(changedEntities, null, null); + } + writeBulkChangeEvents(changed, eventType, userName, phasePrefix + "ChangeEvents"); + } + + private void writeBulkVersionHistory(List changed, String phasePrefix) { + try (var ignored = phase(phasePrefix + "VersionHistory")) { + List historyIds = new ArrayList<>(); + List historyExtensions = new ArrayList<>(); + List historyJsons = new ArrayList<>(); + for (EntityUpdater u : changed) { + if (u.isVersionChanged()) { + historyIds.add(u.getOriginal().getId()); + historyExtensions.add( + EntityUtil.getVersionExtension(entityType, u.getOriginal().getVersion())); + historyJsons.add(JsonUtils.pojoToJson(u.getOriginal())); + } + } + if (!historyIds.isEmpty()) { + daoCollection + .entityExtensionDAO() + .insertMany(historyIds, historyExtensions, entityType, historyJsons); + } + } + } + + private void writeBulkChangeEvents( + List changed, EventType eventType, String userName, String phaseName) { + try (var ignored = phase(phaseName)) { + List changeEventJsons = new ArrayList<>(); + for (EntityUpdater u : changed) { + buildChangeEventJsonForBulkOperation(u.getUpdated(), eventType, userName) + .ifPresent(changeEventJsons::add); } + insertChangeEventsBatch(changeEventJsons); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/FeedRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/FeedRepository.java index 6b338a40379a..124eaca672b7 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/FeedRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/FeedRepository.java @@ -51,10 +51,12 @@ import java.util.ArrayList; import java.util.Comparator; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.Set; import java.util.UUID; import java.util.stream.Collectors; import lombok.Getter; @@ -760,49 +762,95 @@ public void deleteThreadInternal(UUID id) { dao.feedDAO().delete(getLegacyThreadTableName(), id); } + // Keep IN-list expansions well under MySQL's max_allowed_packet budget and within + // PostgreSQL's bind-parameter ceiling. 500 also matches the existing + // EntityRepository.RELATION_DELETE_BATCH_SIZE used for the same reason on the + // relationship side. Smaller than EntityDAO.MAX_IN_LIST_CHUNK_SIZE because the + // feed cleanup path issues three IN-list statements per chunk (relationships, + // field_relationship, thread_entity) and each has its own packet/parameter budget. + private static final int FEED_IN_BATCH_SIZE = 500; + @Transaction public int deleteThreadsInBatch(List threadUUIDs) { if (CommonUtil.nullOrEmpty(threadUUIDs)) return 0; List threadIds = threadUUIDs.stream().map(UUID::toString).toList(); + int deleted = 0; + for (int i = 0; i < threadIds.size(); i += FEED_IN_BATCH_SIZE) { + List chunk = threadIds.subList(i, Math.min(i + FEED_IN_BATCH_SIZE, threadIds.size())); - // Delete all the relationships to other entities - dao.relationshipDAO().deleteAllByThreadIds(threadIds, Entity.THREAD); + // Delete all the relationships to other entities + dao.relationshipDAO().deleteAllByThreadIds(chunk, Entity.THREAD); - // Delete all the field relationships to other entities - dao.fieldRelationshipDAO().deleteAllByPrefixes(threadIds); + // Delete all the field relationships to other entities + dao.fieldRelationshipDAO().deleteAllByPrefixes(chunk); - // Delete the thread and return the count - return dao.feedDAO().deleteByIds(getLegacyThreadTableName(), threadIds); + // Delete the threads in this chunk and tally the count + deleted += dao.feedDAO().deleteByIds(getLegacyThreadTableName(), chunk); + } + return deleted; } public void deleteByAbout(UUID entityId) { + deleteByAbout(List.of(entityId)); + } + + public void deleteByAbout(List entityIds) { + if (entityIds == null || entityIds.isEmpty()) { + return; + } if (!isLegacyThreadStorageAvailable()) { LOG.debug( - "Skipping legacy feed cleanup for entity {} because thread storage is unavailable", - entityId); + "Skipping legacy feed cleanup for {} entities because thread storage is unavailable", + entityIds.size()); return; } - - List threadIds; - try { - threadIds = - listOrEmpty( - dao.feedDAO().findByEntityId(getLegacyThreadTableName(), entityId.toString())); - } catch (Exception ex) { - LOG.debug( - "Skipping legacy feed cleanup for entity {} because thread storage is unavailable", - entityId, - ex); + List entityIdStrings = entityIds.stream().map(UUID::toString).toList(); + // LinkedHashSet: per-chunk findByEntityIds is already DISTINCT, but accumulating across + // chunks could still see the same id twice if a future caller passes an entityIds list + // with duplicates. Dedup once here so deleteThreadsInBatch's downstream chunking (3 + // IN-list DELETEs per 500-id chunk) doesn't waste budget on redundant rows. Linked + // ordering for deterministic logs / replay. + Set threadIds = new LinkedHashSet<>(); + for (int i = 0; i < entityIdStrings.size(); i += FEED_IN_BATCH_SIZE) { + List chunk = + entityIdStrings.subList(i, Math.min(i + FEED_IN_BATCH_SIZE, entityIdStrings.size())); + try { + threadIds.addAll( + listOrEmpty(dao.feedDAO().findByEntityIds(getLegacyThreadTableName(), chunk))); + } catch (Exception ex) { + LOG.debug( + "Skipping legacy feed cleanup for chunk of {} entities (offset {}) because thread storage is unavailable", + chunk.size(), + i, + ex); + } + } + if (threadIds.isEmpty()) { return; } + // Keep legacy feed cleanup best-effort: a malformed thread id or a DAO failure + // here must not blow up the caller's hard-delete @Transaction. Parse defensively + // (skip + log malformed ids) and swallow batch-delete failures. + List threadUuids = new ArrayList<>(threadIds.size()); for (String threadId : threadIds) { try { - deleteThreadInternal(UUID.fromString(threadId)); - } catch (Exception ex) { - // Continue deletion + threadUuids.add(UUID.fromString(threadId)); + } catch (IllegalArgumentException ex) { + LOG.warn("Skipping malformed legacy thread id {} during feed cleanup", threadId); } } + if (threadUuids.isEmpty()) { + return; + } + try { + deleteThreadsInBatch(threadUuids); + } catch (Exception ex) { + LOG.warn( + "Legacy feed cleanup failed for {} threads; continuing entity delete", + threadUuids.size(), + ex); + } } private boolean isLegacyThreadStorageAvailable() { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java index b7e21d374e50..c83ab69609ec 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ListFilter.java @@ -977,10 +977,49 @@ public static String escapeApostrophe(String name) { return name.replace("'", "''"); } + /** + * Defence-in-depth: when a value is embedded inside a single-quoted SQL string literal, + * escape backslashes before apostrophes (MySQL treats {@code \} as a string-literal escape + * by default, and Postgres does too when {@code standard_conforming_strings = off}). Run + * this BEFORE {@link #escapeApostrophe} so the {@code \\} we just inserted isn't itself + * re-doubled. + */ + public static String escapeBackslashAndApostrophe(String name) { + return escapeApostrophe(name.replace("\\", "\\\\")); + } + + /** + * Escape a string for use as the replacement argument to MySQL's + * {@code REGEXP_REPLACE}. Two layers of escaping are needed: + *

    + *
  1. Regex replacement layer: {@code REGEXP_REPLACE} treats {@code \} as the start of a + * backreference / escape sequence (e.g. {@code \1} resolves to capture group 1). + * Each literal backslash in the input needs to become {@code \\} for the regex + * engine to emit a single {@code \}.
    2. + *
  2. SQL string-literal layer: the regex-escaped value is then embedded inside a + * single-quoted SQL string, so each remaining {@code \} doubles again + * ({@code \\} → {@code \\\\}) and apostrophes double ({@code '} → {@code ''}).
    3. + *
+ * Net effect: one input backslash → four backslashes in the SQL statement text, which + * the SQL parser folds to two backslashes for the regex engine, which the regex engine + * folds to one literal backslash in the replacement output. Apostrophes just double + * once (regex replacement doesn't reserve apostrophes, only the SQL layer does). + * + *

Compose with {@link #escapeApostrophe} rather than {@link #escapeBackslashAndApostrophe} + * for the second pass — applying {@code escapeBackslashAndApostrophe} twice would + * re-escape the apostrophes we already doubled. + */ + public static String escapeForMySqlRegexReplacement(String name) { + // Step 1: double backslashes for the regex replacement layer. + String regexEscaped = name.replace("\\", "\\\\"); + // Step 2: double backslashes (again) + apostrophes for the SQL string-literal layer. + return escapeBackslashAndApostrophe(regexEscaped); + } + public static String escape(String name) { // Escape string to be using in LIKE clause // "'" is used for indicated start and end of the string. Use "''" to escape it. - name = escapeApostrophe(name); + name = escapeBackslashAndApostrophe(name); // "_" is a wildcard and looks for any single character. Add "\\" in front of it to escape it return name.replaceAll("_", "\\\\_"); } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java index 38ccd4687bc0..e8aeb4908ddf 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/TestCaseRepository.java @@ -878,10 +878,17 @@ public void clearParentCache() { } private void updateTestSuite(TestCase testCase) { - var testSuiteRepository = (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE); - TestSuite testSuite = Entity.getEntity(testCase.getTestSuite(), "*", ALL); - var original = TestSuiteRepository.copyTestSuite(testSuite); - testSuiteRepository.postUpdate(original, testSuite); + if (testCase.getTestSuite() != null) { + try { + var testSuiteRepository = + (TestSuiteRepository) Entity.getEntityRepository(Entity.TEST_SUITE); + TestSuite testSuite = Entity.getEntity(testCase.getTestSuite(), "*", ALL); + var original = TestSuiteRepository.copyTestSuite(testSuite); + testSuiteRepository.postUpdate(original, testSuite); + } catch (EntityNotFoundException ignored) { + // TestSuite already deleted as part of the same cascade — nothing to update. + } + } } private void updateLogicalTestSuite(UUID testSuiteId) { diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/EntityResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/EntityResource.java index c8c5a48af470..57f62d757e51 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/EntityResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/EntityResource.java @@ -72,7 +72,9 @@ import org.openmetadata.service.OpenMetadataApplicationConfig; import org.openmetadata.service.cache.CacheBundle; import org.openmetadata.service.cache.CacheProvider; +import org.openmetadata.service.exception.BadRequestException; import org.openmetadata.service.exception.CatalogExceptionMessage; +import org.openmetadata.service.exception.EntityNotFoundException; import org.openmetadata.service.jdbi3.EntityRepository; import org.openmetadata.service.jdbi3.ListFilter; import org.openmetadata.service.limits.Limits; @@ -103,6 +105,7 @@ import org.openmetadata.service.util.RestUtil.DeleteResponse; import org.openmetadata.service.util.RestUtil.PatchResponse; import org.openmetadata.service.util.RestUtil.PutResponse; +import org.openmetadata.service.util.RestoreEntityResponse; import org.openmetadata.service.util.ValidatorUtil; import org.openmetadata.service.util.WebsocketNotificationHandler; @@ -771,11 +774,32 @@ public Response deleteByName( } public Response restoreEntity(UriInfo uriInfo, SecurityContext securityContext, UUID id) { + // Read ?async=true off uriInfo so subclass resources that haven't (yet) declared the + // QueryParam still honor the async contract. Lets SDK callers opt into async restore + // universally regardless of which Resource subclass forwarded the parameter. + boolean asyncFromQuery = + uriInfo != null && Boolean.parseBoolean(uriInfo.getQueryParameters().getFirst("async")); + return restoreEntity(uriInfo, securityContext, id, asyncFromQuery); + } + + public Response restoreEntity( + UriInfo uriInfo, SecurityContext securityContext, UUID id, boolean async) { + if (async) { + return restoreEntityAsync(uriInfo, securityContext, id); + } OperationContext operationContext = new OperationContext(entityType, MetadataOperation.EDIT_ALL); authorizer.authorize(securityContext, operationContext, getResourceContextById(id)); PutResponse response = repository.restoreEntity(securityContext.getUserPrincipal().getName(), id); + if (response == null) { + // EntityRepository.restoreEntity now calls find(id, Include.ALL) up front, so a truly + // missing id has already propagated EntityNotFoundException (→ 404) before we got + // here. A null response can only mean "entity exists but is not in DELETED state" — + // map that to 400. + throw new BadRequestException( + String.format("Entity %s:%s is not in deleted state", entityType, id)); + } repository.restoreFromSearch(response.getEntity()); addHref(uriInfo, response.getEntity()); LOG.info( @@ -785,6 +809,109 @@ public Response restoreEntity(UriInfo uriInfo, SecurityContext securityContext, return response.toResponse(); } + /** + * Async restore variant. Returns 202 Accepted with a job ID and runs the restore on the + * shared async executor. The caller can subscribe to + * {@link org.openmetadata.service.socket.WebSocketManager#RESTORE_ENTITY_CHANNEL} to be + * notified when the restore completes or fails. Used to avoid proxy / ALB idle timeouts on + * large hierarchies (issue #4003). + */ + public Response restoreEntityAsync(UriInfo uriInfo, SecurityContext securityContext, UUID id) { + OperationContext operationContext = + new OperationContext(entityType, MetadataOperation.EDIT_ALL); + authorizer.authorize(securityContext, operationContext, getResourceContextById(id)); + // Cheap pre-check so we return a synchronous error instead of 202 + delayed WebSocket + // FAILED for a request that can't succeed. Distinguish the two failure modes that the + // raw EntityNotFoundException would conflate: 404 if the entity truly doesn't exist + // (Include.ALL still finds nothing) and 400 if it exists but is already restored. + // Capturing the entity here also yields a meaningful name for any later FAILED + // notification. + T preCheck; + try { + preCheck = repository.find(id, Include.DELETED); + } catch (EntityNotFoundException notDeleted) { + // Probe with Include.ALL to distinguish 404-missing from 400-not-deleted. Narrow + // catch so unrelated failures (DB connectivity, auth) propagate naturally rather + // than being mis-mapped to 400 "not in deleted state". + boolean entityExists; + try { + repository.find(id, Include.ALL); + entityExists = true; + } catch (EntityNotFoundException missing) { + entityExists = false; + } + if (entityExists) { + throw new BadRequestException( + String.format("Entity %s:%s is not in deleted state", entityType, id)); + } + throw notDeleted; + } + String entityName = preCheck.getName() != null ? preCheck.getName() : id.toString(); + String jobId = UUID.randomUUID().toString(); + String userName = securityContext.getUserPrincipal().getName(); + // Resolve the WebSocket user id on the request thread, while the SecurityContext is + // still valid. JAX-RS may invalidate request-scoped state once the 202 response is + // returned, so we cannot rely on securityContext.getUserPrincipal() inside the lambda. + UUID notifyUserId = WebsocketNotificationHandler.resolveUserId(securityContext); + ExecutorService executorService = AsyncService.getInstance().getExecutorService(); + // Intentionally don't capture uriInfo in the lambda — same request-scope concern. The + // WebSocket notification only needs name/status, not HREFs. + executorService.submit( + RequestLatencyContext.wrapWithContext( + () -> { + try { + PutResponse response = repository.restoreEntity(userName, id); + if (response == null) { + // Pre-check saw the entity in DELETED state; a null response now means a + // concurrent restore won the race. Treat as idempotent success — the + // operator's request is satisfied. If the entity has since been hard- + // deleted, surface that as a real failure. + handleAlreadyRestored(jobId, id, entityName, notifyUserId); + return; + } + repository.restoreFromSearch(response.getEntity()); + LOG.info( + "[AsyncRestore] Restored {}:{} (jobId={})", + Entity.getEntityTypeFromObject(response.getEntity()), + response.getEntity().getId(), + jobId); + WebsocketNotificationHandler.sendRestoreOperationCompleteNotification( + jobId, notifyUserId, response.getEntity()); + } catch (Exception e) { + LOG.error( + "[AsyncRestore] Failed to restore {}:{} (name={})", + entityType, + id, + entityName, + e); + WebsocketNotificationHandler.sendRestoreOperationFailedNotification( + jobId, + notifyUserId, + entityName, + e.getMessage() == null ? e.toString() : e.getMessage()); + } + })); + RestoreEntityResponse response = + new RestoreEntityResponse(jobId, "Restore initiated successfully."); + return Response.accepted().entity(response).type(MediaType.APPLICATION_JSON).build(); + } + + private void handleAlreadyRestored(String jobId, UUID id, String entityName, UUID notifyUserId) { + try { + T restored = repository.find(id, Include.NON_DELETED); + LOG.info( + "[AsyncRestore] {} {} was already restored by another request (jobId={})", + entityType, + id, + jobId); + WebsocketNotificationHandler.sendRestoreOperationCompleteNotification( + jobId, notifyUserId, restored); + } catch (EntityNotFoundException missing) { + WebsocketNotificationHandler.sendRestoreOperationFailedNotification( + jobId, notifyUserId, entityName, "Entity was hard-deleted before restore"); + } + } + public Response exportCsvInternalAsync( SecurityContext securityContext, String name, boolean recursive) { OperationContext operationContext = diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dashboards/DashboardResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/dashboards/DashboardResource.java index 3ac05657402a..0edbda411716 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/dashboards/DashboardResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/dashboards/DashboardResource.java @@ -612,7 +612,9 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted dashboard", - description = "Restore a soft deleted dashboard.", + description = + "Restore a soft deleted dashboard. Pass async=true to run the restore in the background" + + " and receive a 202 Accepted response with a job id.", responses = { @ApiResponse( responseCode = "200", @@ -620,12 +622,26 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = Dashboard.class))) + schema = @Schema(implementation = Dashboard.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreDashboard( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseResource.java index 200170085ec4..a35496332d96 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseResource.java @@ -751,7 +751,11 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted Database.", - description = "Restore a soft deleted Database.", + description = + "Restore a soft deleted Database. Pass async=true to run the restore in the" + + " background and receive a 202 Accepted response with a job id; useful for" + + " hierarchies large enough that the synchronous response would exceed proxy" + + " idle timeouts.", responses = { @ApiResponse( responseCode = "200", @@ -759,13 +763,27 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = Database.class))) + schema = @Schema(implementation = Database.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreDatabase( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @PUT diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseSchemaResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseSchemaResource.java index 7be0e5fa9441..42e14d579c07 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseSchemaResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/DatabaseSchemaResource.java @@ -812,7 +812,10 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted database schema.", - description = "Restore a soft deleted database schema.", + description = + "Restore a soft deleted database schema. Pass async=true to run the restore in the" + + " background and receive a 202 Accepted response with a job id; useful when the" + + " schema contains thousands of tables.", responses = { @ApiResponse( responseCode = "200", @@ -820,13 +823,27 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = DatabaseSchema.class))) + schema = @Schema(implementation = DatabaseSchema.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreDatabaseSchema( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @PUT diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/StoredProcedureResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/StoredProcedureResource.java index 760da9de6b55..ee00edcd1d86 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/StoredProcedureResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/StoredProcedureResource.java @@ -600,7 +600,9 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted stored procedure.", - description = "Restore a soft deleted stored procedure.", + description = + "Restore a soft deleted stored procedure. Pass async=true to run the restore in the" + + " background and receive a 202 Accepted response with a job id.", responses = { @ApiResponse( responseCode = "200", @@ -608,12 +610,26 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = StoredProcedure.class))) + schema = @Schema(implementation = StoredProcedure.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreStoredProcedure( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/TableResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/TableResource.java index fea914d1726f..cef472f41000 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/TableResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/databases/TableResource.java @@ -783,7 +783,9 @@ public Response deleteByFqn( @Operation( operationId = "restore", summary = "Restore a soft deleted table", - description = "Restore a soft deleted table.", + description = + "Restore a soft deleted table. Pass async=true to run the restore in the background" + + " and receive a 202 Accepted response with a job id.", responses = { @ApiResponse( responseCode = "200", @@ -791,13 +793,27 @@ public Response deleteByFqn( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = Table.class))) + schema = @Schema(implementation = Table.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreTable( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @PUT diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/datamodels/DashboardDataModelResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/datamodels/DashboardDataModelResource.java index 8a69e00aa088..e27d43f34839 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/datamodels/DashboardDataModelResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/datamodels/DashboardDataModelResource.java @@ -625,7 +625,9 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted data model.", - description = "Restore a soft deleted data model.", + description = + "Restore a soft deleted data model. Pass async=true to run the restore in the" + + " background and receive a 202 Accepted response with a job id.", responses = { @ApiResponse( responseCode = "200", @@ -633,13 +635,27 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = DashboardDataModel.class))) + schema = @Schema(implementation = DashboardDataModel.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreDataModel( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @GET diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/services/database/DatabaseServiceResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/services/database/DatabaseServiceResource.java index 394a145c493f..fc93f9384842 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/services/database/DatabaseServiceResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/services/database/DatabaseServiceResource.java @@ -769,7 +769,10 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted database service", - description = "Restore a soft deleted database service.", + description = + "Restore a soft deleted database service. Pass async=true to run the restore in the" + + " background and receive a 202 Accepted response with a job id; strongly" + + " recommended for services that contain many databases / schemas / tables.", responses = { @ApiResponse( responseCode = "200", @@ -777,13 +780,27 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = DatabaseService.class))) + schema = @Schema(implementation = DatabaseService.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreDatabaseService( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @Override diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/storages/ContainerResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/storages/ContainerResource.java index 7253558c5979..5f27763ed53f 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/storages/ContainerResource.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/storages/ContainerResource.java @@ -621,7 +621,10 @@ public Response delete( @Operation( operationId = "restore", summary = "Restore a soft deleted Container.", - description = "Restore a soft deleted Container.", + description = + "Restore a soft deleted Container. Pass async=true to run the restore in the background" + + " and receive a 202 Accepted response with a job id; useful for deep container" + + " hierarchies.", responses = { @ApiResponse( responseCode = "200", @@ -629,13 +632,27 @@ public Response delete( content = @Content( mediaType = "application/json", - schema = @Schema(implementation = Container.class))) + schema = @Schema(implementation = Container.class))), + @ApiResponse( + responseCode = "202", + description = "Async restore started. Track completion via the jobId.", + content = + @Content( + mediaType = "application/json", + schema = + @Schema( + implementation = + org.openmetadata.service.util.RestoreEntityResponse.class))) }) public Response restoreContainer( @Context UriInfo uriInfo, @Context SecurityContext securityContext, + @Parameter(description = "Run the restore asynchronously. (Default = `false`)") + @QueryParam("async") + @DefaultValue("false") + boolean async, @Valid RestoreEntity restore) { - return restoreEntity(uriInfo, securityContext, restore.getId()); + return restoreEntity(uriInfo, securityContext, restore.getId(), async); } @PUT diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/socket/WebSocketManager.java b/openmetadata-service/src/main/java/org/openmetadata/service/socket/WebSocketManager.java index 57ffaaa46a62..387b19282876 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/socket/WebSocketManager.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/socket/WebSocketManager.java @@ -36,6 +36,7 @@ public class WebSocketManager { public static final String BULK_ASSETS_CHANNEL = "bulkAssetsChannel"; public static final String DELETE_ENTITY_CHANNEL = "deleteEntityChannel"; + public static final String RESTORE_ENTITY_CHANNEL = "restoreEntityChannel"; public static final String MOVE_GLOSSARY_TERM_CHANNEL = "moveGlossaryTermChannel"; public static final String RDF_INDEX_JOB_BROADCAST_CHANNEL = "rdfIndexJobStatus"; public static final String CHART_DATA_STREAM_CHANNEL = "chartDataStream"; diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/AsyncService.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/AsyncService.java index 8e6d74405b1a..b01ee4607c94 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/AsyncService.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/AsyncService.java @@ -1,68 +1,42 @@ package org.openmetadata.service.util; -import java.util.List; -import java.util.concurrent.AbstractExecutorService; import java.util.concurrent.Callable; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import java.util.function.Supplier; -import lombok.Getter; import lombok.extern.slf4j.Slf4j; -import org.openmetadata.service.OpenMetadataApplicationConfigHolder; +/** + * Single virtual-thread executor for all server-side async dispatch (CSV export/import, + * bulk asset ops, async delete/restore). + * + *

Back-pressure is intentionally not enforced here. The old semaphore-based + * bounded wrapper was fighting Project Loom — virtual threads scale to millions and are + * basically free, while the real bottleneck under load is the JDBI connection pool. Letting + * tasks queue on connection acquisition (with the pool's own timeout) is both simpler and + * more accurate than guessing at "how many concurrent tasks ≈ connection pool capacity". + * + *

If a future use case genuinely needs admission control, it should live at the caller + * boundary (e.g., a token bucket per user, or a per-operation queue with rejection) rather + * than at this shared executor. + */ @Slf4j public class AsyncService { private static AsyncService instance; private final ExecutorService executorService; - private final Semaphore concurrencyLimiter; - @Getter private final int maxConcurrency; private static final int DEFAULT_MAX_RETRIES = 3; private static final long DEFAULT_INITIAL_RETRY_DELAY_MS = 1000; private static final long DEFAULT_OPERATION_TIMEOUT_SECONDS = 60; - private static final long SHUTDOWN_TIMEOUT_SECONDS = 30; private AsyncService() { - maxConcurrency = resolveMaxConcurrency(); - concurrencyLimiter = new Semaphore(maxConcurrency); executorService = - new BoundedExecutorService( - Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("om-async-", 0).factory()), - concurrencyLimiter); - LOG.info("AsyncService initialized with max concurrency: {}", maxConcurrency); - } - - private static int resolveMaxConcurrency() { - String env = System.getenv("ASYNC_SERVICE_MAX_CONCURRENCY"); - if (env != null) { - try { - int value = Integer.parseInt(env.trim()); - if (value > 0) { - return value; - } - } catch (NumberFormatException ignored) { - } - } - int cpuBudget = Runtime.getRuntime().availableProcessors() * 2; - try { - if (OpenMetadataApplicationConfigHolder.isInitialized()) { - int poolSize = - OpenMetadataApplicationConfigHolder.getInstance().getDataSourceFactory().getMaxSize(); - if (poolSize > 0) { - return Math.max(4, Math.min(cpuBudget, poolSize / 3)); - } - } - } catch (Exception e) { - LOG.warn( - "Could not determine database pool size, using CPU-based concurrency budget: {}", - e.getMessage()); - } - return Math.max(4, cpuBudget); + Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("om-async-", 0).factory()); + LOG.info("AsyncService initialized (virtual-thread-per-task executor)"); } public static synchronized AsyncService getInstance() { @@ -95,7 +69,7 @@ public CompletableFuture submit(Callable task) { } public void shutdown() { - LOG.info("Shutting down AsyncService executor (max concurrency: {})", maxConcurrency); + LOG.info("Shutting down AsyncService executor"); executorService.shutdown(); try { if (!executorService.awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) { @@ -233,63 +207,4 @@ private static T executeWithRetry( throw new RuntimeException( String.format("Failed to %s %s", operationName.toLowerCase(), context), lastException); } - - /** - * ExecutorService wrapper that enforces concurrency limits via a semaphore. Every task submitted - * through any method (execute, submit, invokeAll, invokeAny) acquires a permit before running and - * releases it on completion. This ensures ALL callers — including those using getExecutorService() - * directly — are bounded. - */ - private static class BoundedExecutorService extends AbstractExecutorService { - private final ExecutorService delegate; - private final Semaphore semaphore; - - BoundedExecutorService(ExecutorService delegate, Semaphore semaphore) { - this.delegate = delegate; - this.semaphore = semaphore; - } - - @Override - public void execute(Runnable command) { - delegate.execute( - () -> { - try { - semaphore.acquire(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw new RuntimeException("Interrupted waiting for concurrency permit", e); - } - try { - command.run(); - } finally { - semaphore.release(); - } - }); - } - - @Override - public void shutdown() { - delegate.shutdown(); - } - - @Override - public List shutdownNow() { - return delegate.shutdownNow(); - } - - @Override - public boolean isShutdown() { - return delegate.isShutdown(); - } - - @Override - public boolean isTerminated() { - return delegate.isTerminated(); - } - - @Override - public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException { - return delegate.awaitTermination(timeout, unit); - } - } } diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityMessage.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityMessage.java new file mode 100644 index 000000000000..7e3d8216fe77 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityMessage.java @@ -0,0 +1,32 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.service.util; + +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@NoArgsConstructor +public class RestoreEntityMessage { + @Getter @Setter private String jobId; + @Getter @Setter private String status; + @Getter @Setter private String entityName; + @Getter @Setter private String error; + + public RestoreEntityMessage(String jobId, String status, String entityName, String error) { + this.jobId = jobId; + this.status = status; + this.entityName = entityName; + this.error = error; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityResponse.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityResponse.java new file mode 100644 index 000000000000..b55519478cf4 --- /dev/null +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/RestoreEntityResponse.java @@ -0,0 +1,34 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.service.util; + +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +/** + * Response shape for an async restore request. Returned with HTTP 202 when a client passes + * {@code async=true} to the restore endpoint. The {@code jobId} can be used to correlate + * subsequent WebSocket notifications on + * {@link org.openmetadata.service.socket.WebSocketManager#RESTORE_ENTITY_CHANNEL}. + */ +@NoArgsConstructor +public class RestoreEntityResponse { + @Getter @Setter private String jobId; + @Getter @Setter private String message; + + public RestoreEntityResponse(String jobId, String message) { + this.jobId = jobId; + this.message = message; + } +} diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/WebsocketNotificationHandler.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/WebsocketNotificationHandler.java index 0a765365bc58..569d4f83a024 100644 --- a/openmetadata-service/src/main/java/org/openmetadata/service/util/WebsocketNotificationHandler.java +++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/WebsocketNotificationHandler.java @@ -420,6 +420,48 @@ public static void sendDeleteOperationFailedNotification( } } + public static void sendRestoreOperationCompleteNotification( + String jobId, UUID userId, EntityInterface entity) { + RestoreEntityMessage message = + new RestoreEntityMessage(jobId, "COMPLETED", entity.getName(), null); + String jsonMessage = JsonUtils.pojoToJson(message); + LOG.info( + "[AsyncRestore] Restore operation completed - jobId: {}, userId: {}, entity: {}", + jobId, + userId, + entity.getName()); + if (userId != null) { + WebSocketManager.getInstance() + .sendToOne(userId, WebSocketManager.RESTORE_ENTITY_CHANNEL, jsonMessage); + } + } + + public static void sendRestoreOperationFailedNotification( + String jobId, UUID userId, String entityName, String error) { + RestoreEntityMessage message = new RestoreEntityMessage(jobId, "FAILED", entityName, error); + String jsonMessage = JsonUtils.pojoToJson(message); + LOG.error( + "[AsyncRestore] Restore operation failed - jobId: {}, userId: {}, entity: {}, error: {}", + jobId, + userId, + entityName, + error); + if (userId != null) { + WebSocketManager.getInstance() + .sendToOne(userId, WebSocketManager.RESTORE_ENTITY_CHANNEL, jsonMessage); + } + } + + /** + * Resolve the WebSocket user id for the given security context. Call this on the + * request thread (i.e., before submitting an async task) so the lookup runs while the + * SecurityContext is still valid — JAX-RS may invalidate request-scoped state after the + * response returns. + */ + public static UUID resolveUserId(SecurityContext securityContext) { + return getUserIdFromSecurityContext(securityContext); + } + public static void sendMoveOperationCompleteNotification( String jobId, SecurityContext securityContext, EntityInterface entity) { MoveGlossaryTermMessage message = diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/EntityRepositoryRestoreTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/EntityRepositoryRestoreTest.java new file mode 100644 index 000000000000..6b368ed810d2 --- /dev/null +++ b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/EntityRepositoryRestoreTest.java @@ -0,0 +1,493 @@ +/* + * Copyright 2026 Collate + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * http://www.apache.org/licenses/LICENSE-2.0 + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.openmetadata.service.jdbi3; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedStatic; +import org.openmetadata.schema.entity.data.Pipeline; +import org.openmetadata.schema.type.Include; +import org.openmetadata.schema.type.Relationship; +import org.openmetadata.service.Entity; +import org.openmetadata.service.util.EntityUtil.Fields; +import org.openmetadata.service.util.EntityUtil.RelationIncludes; + +/** + * Unit tests for the iterative bulk restore + bulk soft-delete + bulk hard-delete paths + * introduced for issue #4003. Verifies the dispatch shape that's testable without spinning + * up the full bulk write path: + * + *

    + *
  • {@link EntityRepository#restoreChildren(UUID, String)} groups CONTAINS + PARENT_OF + * children by entity type and dispatches a single {@link + * EntityRepository#bulkRestoreSubtree(List, String)} call per type (instead of N + * recursive {@code Entity.restoreEntity} calls). The relation set matches + * {@code deleteChildren} so a Team / KnowledgePage / Classification hierarchy is + * restored the same way it was cascade-soft-deleted. + *
  • {@link EntityRepository#deleteChildren(List, boolean, String)} with + * {@code hardDelete=false} dispatches one {@link EntityRepository#bulkSoftDeleteSubtree( + * List, String)} call per type and with {@code hardDelete=true} dispatches one + * {@link EntityRepository#bulkHardDeleteSubtree(List, String)} call per type. + *
  • All three bulk methods bail out cleanly on null / empty inputs. + *
  • All three bulk methods issue a single batched {@code findToBatchAllTypes} per tree + * level that walks both {@code CONTAINS} and {@code PARENT_OF} so Glossary / Team / + * recursive-Container descendants stop silently slipping past the cascade. + *
  • The per-entity {@code *AdditionalChildren} hooks fire even on the "entities present + * but none need flipping" branch (so a re-entered cascade can reconcile HAS-related + * descendants), and {@code hardDeleteAdditionalChildren} + {@code + * bulkEntitySpecificCleanup} fire on the full bulk hard-delete path with the expected + * per-entity / per-batch counts. + *
+ * + * The full bulk DB-write path (version history, updateMany, change events, entity row + * deletes) is exercised in {@code RestoreHierarchyIT}, which runs against a real Docker + * stack. + */ +class EntityRepositoryRestoreTest { + + private static final List SUBTREE_RELATIONS = + List.of(Relationship.CONTAINS.ordinal(), Relationship.PARENT_OF.ordinal()); + + private CollectionDAO daoCollection; + private CollectionDAO.EntityRelationshipDAO relationshipDAO; + private CollectionDAO.PipelineDAO pipelineDAO; + + private static class CountingPipelineRepo extends EntityRepository { + int restoreAdditionalChildrenCalls = 0; + int softDeleteAdditionalChildrenCalls = 0; + int hardDeleteAdditionalChildrenCalls = 0; + int bulkEntitySpecificCleanupCalls = 0; + final Set bulkRestoreInvokedWith = new HashSet<>(); + final Set bulkSoftDeleteInvokedWith = new HashSet<>(); + final Set bulkHardDeleteInvokedWith = new HashSet<>(); + + CountingPipelineRepo(CollectionDAO.PipelineDAO dao) { + super("pipelines", Entity.PIPELINE, Pipeline.class, dao, "", ""); + } + + @Override + protected void setFields(Pipeline entity, Fields fields, RelationIncludes r) {} + + @Override + protected void clearFields(Pipeline entity, Fields fields) {} + + @Override + protected void prepare(Pipeline entity, boolean update) {} + + @Override + protected void storeEntity(Pipeline entity, boolean update) {} + + @Override + protected void storeRelationships(Pipeline entity) {} + + @Override + protected void restoreAdditionalChildren(UUID id, String updatedBy) { + restoreAdditionalChildrenCalls++; + bulkRestoreInvokedWith.add(id); + } + + @Override + protected void softDeleteAdditionalChildren(UUID id, String updatedBy) { + softDeleteAdditionalChildrenCalls++; + bulkSoftDeleteInvokedWith.add(id); + } + + @Override + protected void hardDeleteAdditionalChildren(UUID id, String updatedBy) { + hardDeleteAdditionalChildrenCalls++; + bulkHardDeleteInvokedWith.add(id); + } + + @Override + protected void bulkEntitySpecificCleanup(List entities) { + bulkEntitySpecificCleanupCalls++; + } + } + + @BeforeEach + void setUp() { + daoCollection = mock(CollectionDAO.class); + relationshipDAO = mock(CollectionDAO.EntityRelationshipDAO.class); + pipelineDAO = mock(CollectionDAO.PipelineDAO.class); + when(daoCollection.relationshipDAO()).thenReturn(relationshipDAO); + Entity.setCollectionDAO(daoCollection); + } + + @AfterEach + void tearDown() { + Entity.setCollectionDAO(null); + } + + @Test + void restoreChildren_withNoChildren_isNoOp() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID parentId = UUID.randomUUID(); + when(relationshipDAO.findTo(eq(parentId), eq(Entity.PIPELINE), eq(SUBTREE_RELATIONS))) + .thenReturn(List.of()); + + repo.restoreChildren(parentId, "user"); + + verify(relationshipDAO).findTo(eq(parentId), eq(Entity.PIPELINE), eq(SUBTREE_RELATIONS)); + assertEquals(0, repo.restoreAdditionalChildrenCalls); + } + + @Test + void restoreChildren_groupsByTypeAndDispatchesOnceEach() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID parentId = UUID.randomUUID(); + + UUID schemaA = UUID.randomUUID(); + UUID schemaB = UUID.randomUUID(); + UUID procA = UUID.randomUUID(); + + List children = new ArrayList<>(); + children.add(record(schemaA, Entity.DATABASE_SCHEMA)); + children.add(record(schemaB, Entity.DATABASE_SCHEMA)); + children.add(record(procA, Entity.STORED_PROCEDURE)); + when(relationshipDAO.findTo(eq(parentId), eq(Entity.PIPELINE), eq(SUBTREE_RELATIONS))) + .thenReturn(children); + + EntityRepository schemaRepo = mock(EntityRepository.class); + EntityRepository procRepo = mock(EntityRepository.class); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock + .when(() -> Entity.getEntityRepository(Entity.DATABASE_SCHEMA)) + .thenReturn(schemaRepo); + entityMock + .when(() -> Entity.getEntityRepository(Entity.STORED_PROCEDURE)) + .thenReturn(procRepo); + + repo.restoreChildren(parentId, "user"); + } + + ArgumentCaptor> schemaIds = captureUuidList(); + verify(schemaRepo, times(1)).bulkRestoreSubtree(schemaIds.capture(), eq("user")); + assertEquals(2, schemaIds.getValue().size()); + assertTrue(schemaIds.getValue().contains(schemaA)); + assertTrue(schemaIds.getValue().contains(schemaB)); + + ArgumentCaptor> procIds = captureUuidList(); + verify(procRepo, times(1)).bulkRestoreSubtree(procIds.capture(), eq("user")); + assertEquals(1, procIds.getValue().size()); + assertTrue(procIds.getValue().contains(procA)); + + verify(schemaRepo, never()).restoreEntity(eq("user"), eq(schemaA)); + verify(schemaRepo, never()).restoreEntity(eq("user"), eq(schemaB)); + verify(procRepo, never()).restoreEntity(eq("user"), eq(procA)); + } + + @Test + void bulkRestoreSubtree_emptyOrNullIds_isNoOp() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + + repo.bulkRestoreSubtree(null, "user"); + repo.bulkRestoreSubtree(List.of(), "user"); + + // bulkRestoreSubtree loads with Include.ALL — guard that neither the DELETED nor ALL + // shape is invoked when the input list is empty/null. + verify(pipelineDAO, never()) + .findEntitiesByIds(anyList(), eq(org.openmetadata.schema.type.Include.DELETED)); + verify(pipelineDAO, never()).findEntitiesByIds(anyList(), eq(Include.ALL)); + assertEquals(0, repo.restoreAdditionalChildrenCalls); + } + + @Test + void bulkRestoreSubtree_noEntitiesAtAll_isNoOp() { + // loadForBulk returns an empty list (entity doesn't exist at all): bulk path bails + // before children traversal or hook invocation. + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID id = UUID.randomUUID(); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of()); + + repo.bulkRestoreSubtree(List.of(id), "user"); + + verify(pipelineDAO, atLeastOnce()).findEntitiesByIds(anyList(), eq(Include.ALL)); + assertEquals(0, repo.restoreAdditionalChildrenCalls); + } + + @Test + void bulkRestoreSubtree_entitiesPresentButNoneDeleted_stillRunsAdditionalChildrenHook() { + // loadForBulk returns entities, but none are in DELETED state. Bulk path must skip + // the deferred-store update phase but still call runRestoreAdditionalChildren — a + // re-entered cascade may have HAS-related descendants that need reconciliation. + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID id = UUID.randomUUID(); + Pipeline pa = + new Pipeline().withId(id).withName("a").withFullyQualifiedName("svc.a").withDeleted(false); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of(pa)); + when(relationshipDAO.findToBatchAllTypes(anyList(), eq(SUBTREE_RELATIONS), eq(Include.ALL))) + .thenReturn(List.of()); + + repo.bulkRestoreSubtree(List.of(id), "user"); + + assertEquals(1, repo.restoreAdditionalChildrenCalls); + assertTrue(repo.bulkRestoreInvokedWith.contains(id)); + } + + @Test + void bulkRestoreSubtree_usesBatchedFindToOncePerLevel() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID a = UUID.randomUUID(); + UUID b = UUID.randomUUID(); + Pipeline pa = + new Pipeline().withId(a).withName("a").withFullyQualifiedName("svc.a").withDeleted(true); + Pipeline pb = + new Pipeline().withId(b).withName("b").withFullyQualifiedName("svc.b").withDeleted(true); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of(pa, pb)); + when(relationshipDAO.findToBatchAllTypes(anyList(), eq(SUBTREE_RELATIONS), eq(Include.ALL))) + .thenReturn(List.of()); + + try { + repo.bulkRestoreSubtree(List.of(a, b), "user"); + } catch (Exception ignored) { + // Heavy DB write path requires more wiring than this unit test mocks; we only care + // that the per-level findTo collapse happened before any failure downstream. + } + + ArgumentCaptor> idsCap = captureStringList(); + verify(relationshipDAO, times(1)) + .findToBatchAllTypes(idsCap.capture(), eq(SUBTREE_RELATIONS), eq(Include.ALL)); + assertEquals(2, idsCap.getValue().size()); + assertTrue(idsCap.getValue().contains(a.toString())); + assertTrue(idsCap.getValue().contains(b.toString())); + } + + @Test + void deleteChildren_softDelete_groupsByTypeAndDispatchesToBulkSoftDelete() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + + UUID schemaA = UUID.randomUUID(); + UUID schemaB = UUID.randomUUID(); + UUID procA = UUID.randomUUID(); + + List children = new ArrayList<>(); + children.add(record(schemaA, Entity.DATABASE_SCHEMA)); + children.add(record(schemaB, Entity.DATABASE_SCHEMA)); + children.add(record(procA, Entity.STORED_PROCEDURE)); + + EntityRepository schemaRepo = mock(EntityRepository.class); + EntityRepository procRepo = mock(EntityRepository.class); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock + .when(() -> Entity.getEntityRepository(Entity.DATABASE_SCHEMA)) + .thenReturn(schemaRepo); + entityMock + .when(() -> Entity.getEntityRepository(Entity.STORED_PROCEDURE)) + .thenReturn(procRepo); + + repo.deleteChildren(children, false, "user"); + } + + ArgumentCaptor> schemaIds = captureUuidList(); + verify(schemaRepo, times(1)).bulkSoftDeleteSubtree(schemaIds.capture(), eq("user")); + assertEquals(2, schemaIds.getValue().size()); + assertTrue(schemaIds.getValue().contains(schemaA)); + assertTrue(schemaIds.getValue().contains(schemaB)); + + ArgumentCaptor> procIds = captureUuidList(); + verify(procRepo, times(1)).bulkSoftDeleteSubtree(procIds.capture(), eq("user")); + assertEquals(1, procIds.getValue().size()); + assertTrue(procIds.getValue().contains(procA)); + } + + @Test + void bulkSoftDeleteSubtree_emptyOrNullIds_isNoOp() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + + repo.bulkSoftDeleteSubtree(null, "user"); + repo.bulkSoftDeleteSubtree(List.of(), "user"); + + verify(pipelineDAO, never()).findEntitiesByIds(anyList(), eq(Include.ALL)); + assertEquals(0, repo.softDeleteAdditionalChildrenCalls); + } + + @Test + void bulkSoftDeleteSubtree_usesBatchedFindToOncePerLevel() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID a = UUID.randomUUID(); + UUID b = UUID.randomUUID(); + Pipeline pa = new Pipeline().withId(a).withName("a").withFullyQualifiedName("svc.a"); + Pipeline pb = new Pipeline().withId(b).withName("b").withFullyQualifiedName("svc.b"); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of(pa, pb)); + when(relationshipDAO.findToBatchAllTypes(anyList(), eq(SUBTREE_RELATIONS), eq(Include.ALL))) + .thenReturn(List.of()); + + try { + repo.bulkSoftDeleteSubtree(List.of(a, b), "user"); + } catch (Exception ignored) { + // Heavy DB write path is not mocked; we verify only the per-level findTo collapse. + } + + ArgumentCaptor> idsCap = captureStringList(); + verify(relationshipDAO, times(1)) + .findToBatchAllTypes(idsCap.capture(), eq(SUBTREE_RELATIONS), eq(Include.ALL)); + assertEquals(2, idsCap.getValue().size()); + assertTrue(idsCap.getValue().contains(a.toString())); + assertTrue(idsCap.getValue().contains(b.toString())); + } + + @Test + void deleteChildren_hardDelete_groupsByTypeAndDispatchesToBulkHardDelete() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + + UUID schemaA = UUID.randomUUID(); + UUID schemaB = UUID.randomUUID(); + UUID procA = UUID.randomUUID(); + + List children = new ArrayList<>(); + children.add(record(schemaA, Entity.DATABASE_SCHEMA)); + children.add(record(schemaB, Entity.DATABASE_SCHEMA)); + children.add(record(procA, Entity.STORED_PROCEDURE)); + + EntityRepository schemaRepo = mock(EntityRepository.class); + EntityRepository procRepo = mock(EntityRepository.class); + + try (MockedStatic entityMock = mockStatic(Entity.class)) { + entityMock + .when(() -> Entity.getEntityRepository(Entity.DATABASE_SCHEMA)) + .thenReturn(schemaRepo); + entityMock + .when(() -> Entity.getEntityRepository(Entity.STORED_PROCEDURE)) + .thenReturn(procRepo); + + repo.deleteChildren(children, true, "user"); + } + + ArgumentCaptor> schemaIds = captureUuidList(); + verify(schemaRepo, times(1)).bulkHardDeleteSubtree(schemaIds.capture(), eq("user")); + assertEquals(2, schemaIds.getValue().size()); + assertTrue(schemaIds.getValue().contains(schemaA)); + assertTrue(schemaIds.getValue().contains(schemaB)); + + ArgumentCaptor> procIds = captureUuidList(); + verify(procRepo, times(1)).bulkHardDeleteSubtree(procIds.capture(), eq("user")); + assertEquals(1, procIds.getValue().size()); + assertTrue(procIds.getValue().contains(procA)); + + verify(schemaRepo, never()).bulkSoftDeleteSubtree(anyList(), eq("user")); + verify(procRepo, never()).bulkSoftDeleteSubtree(anyList(), eq("user")); + } + + @Test + void bulkHardDeleteSubtree_emptyOrNullIds_isNoOp() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + + repo.bulkHardDeleteSubtree(null, "user"); + repo.bulkHardDeleteSubtree(List.of(), "user"); + + verify(pipelineDAO, never()).findEntitiesByIds(anyList(), eq(Include.ALL)); + assertEquals(0, repo.hardDeleteAdditionalChildrenCalls); + assertEquals(0, repo.bulkEntitySpecificCleanupCalls); + } + + @Test + void bulkHardDeleteSubtree_usesBatchedFindToOncePerLevel_includingParentOf() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID a = UUID.randomUUID(); + UUID b = UUID.randomUUID(); + Pipeline pa = new Pipeline().withId(a).withName("a").withFullyQualifiedName("svc.a"); + Pipeline pb = new Pipeline().withId(b).withName("b").withFullyQualifiedName("svc.b"); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of(pa, pb)); + when(relationshipDAO.findToBatchAllTypes(anyList(), eq(SUBTREE_RELATIONS), eq(Include.ALL))) + .thenReturn(List.of()); + + try { + repo.bulkHardDeleteSubtree(List.of(a, b), "user"); + } catch (Exception ignored) { + // Heavy DB write path is not mocked; we verify only the per-level findTo collapse and + // hook invocation. + } + + ArgumentCaptor> idsCap = captureStringList(); + verify(relationshipDAO, times(1)) + .findToBatchAllTypes(idsCap.capture(), eq(SUBTREE_RELATIONS), eq(Include.ALL)); + assertEquals(2, idsCap.getValue().size()); + assertTrue(idsCap.getValue().contains(a.toString())); + assertTrue(idsCap.getValue().contains(b.toString())); + } + + @Test + void bulkHardDeleteSubtree_callsBulkEntitySpecificCleanupAndAdditionalChildrenHooks() { + CountingPipelineRepo repo = new CountingPipelineRepo(pipelineDAO); + UUID a = UUID.randomUUID(); + UUID b = UUID.randomUUID(); + Pipeline pa = new Pipeline().withId(a).withName("a").withFullyQualifiedName("svc.a"); + Pipeline pb = new Pipeline().withId(b).withName("b").withFullyQualifiedName("svc.b"); + when(pipelineDAO.findEntitiesByIds(anyList(), eq(Include.ALL))).thenReturn(List.of(pa, pb)); + when(relationshipDAO.findToBatchAllTypes(anyList(), eq(SUBTREE_RELATIONS), eq(Include.ALL))) + .thenReturn(List.of()); + + CollectionDAO.EntityExtensionDAO extensionDAO = mock(CollectionDAO.EntityExtensionDAO.class); + CollectionDAO.FieldRelationshipDAO fieldRelationshipDAO = + mock(CollectionDAO.FieldRelationshipDAO.class); + CollectionDAO.TagUsageDAO tagUsageDAO = mock(CollectionDAO.TagUsageDAO.class); + CollectionDAO.UsageDAO usageDAO = mock(CollectionDAO.UsageDAO.class); + when(daoCollection.entityExtensionDAO()).thenReturn(extensionDAO); + when(daoCollection.fieldRelationshipDAO()).thenReturn(fieldRelationshipDAO); + when(daoCollection.tagUsageDAO()).thenReturn(tagUsageDAO); + when(daoCollection.usageDAO()).thenReturn(usageDAO); + + FeedRepository feedRepository = mock(FeedRepository.class); + try (MockedStatic entityMock = mockStatic(Entity.class, CALLS_REAL_METHODS)) { + entityMock.when(Entity::getFeedRepository).thenReturn(feedRepository); + repo.bulkHardDeleteSubtree(List.of(a, b), "user"); + } + + // bulkEntitySpecificCleanup is invoked once per bulk call with the whole batch. + assertEquals(1, repo.bulkEntitySpecificCleanupCalls); + // hardDeleteAdditionalChildren is invoked once per entity in the batch. + assertEquals(2, repo.hardDeleteAdditionalChildrenCalls); + assertTrue(repo.bulkHardDeleteInvokedWith.contains(a)); + assertTrue(repo.bulkHardDeleteInvokedWith.contains(b)); + // Verify the per-batch relationship + extension cleanup actually ran. + verify(relationshipDAO, times(1)).batchDeleteRelationships(anyList(), eq(Entity.PIPELINE)); + verify(extensionDAO, times(1)).deleteAllBatch(anyList()); + verify(pipelineDAO, times(1)).deleteByIds(anyList()); + } + + private CollectionDAO.EntityRelationshipRecord record(UUID id, String type) { + return CollectionDAO.EntityRelationshipRecord.builder().id(id).type(type).build(); + } + + @SuppressWarnings("unchecked") + private static ArgumentCaptor> captureUuidList() { + return ArgumentCaptor.forClass(List.class); + } + + @SuppressWarnings("unchecked") + private static ArgumentCaptor> captureStringList() { + return ArgumentCaptor.forClass(List.class); + } +} diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java index 06d9d779dc35..50c1937401f6 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/jdbi3/ListFilterTest.java @@ -19,6 +19,81 @@ void test_escapeApostrophe() { assertEquals("a\\_b\\_c\\_d", ListFilter.escape("a_b_c_d")); } + @Test + void test_escapeBackslashAndApostrophe_passesThroughPlainStrings() { + assertEquals("abcd", ListFilter.escapeBackslashAndApostrophe("abcd")); + assertEquals("", ListFilter.escapeBackslashAndApostrophe("")); + } + + @Test + void test_escapeBackslashAndApostrophe_doublesApostrophes() { + // ' → '' for the SQL string-literal layer + assertEquals("a''b", ListFilter.escapeBackslashAndApostrophe("a'b")); + assertEquals("''", ListFilter.escapeBackslashAndApostrophe("'")); + } + + @Test + void test_escapeBackslashAndApostrophe_doublesBackslashesBeforeApostrophes() { + // \ → \\ for the SQL string-literal layer (MySQL default + Postgres legacy mode); + // backslash escape must run BEFORE apostrophe escape so the \\ we just inserted + // is not itself re-doubled by a subsequent pass. + assertEquals("a\\\\b", ListFilter.escapeBackslashAndApostrophe("a\\b")); + assertEquals("\\\\\\\\", ListFilter.escapeBackslashAndApostrophe("\\\\")); + assertEquals("a\\\\''b", ListFilter.escapeBackslashAndApostrophe("a\\'b")); + } + + @Test + void test_escape_alsoDoublesBackslashesViaBackslashAndApostrophe() { + // Regression guard: escape() composes through escapeBackslashAndApostrophe, so a + // literal backslash in the input must come out doubled (defence-in-depth against + // SQL string-literal escape interpretation, on top of the existing LIKE underscore + // escape). + assertEquals("a\\\\b", ListFilter.escape("a\\b")); + assertEquals("a\\\\b\\_c", ListFilter.escape("a\\b_c")); + } + + @Test + void test_escapeForMySqlRegexReplacement_passesThroughPlainStrings() { + assertEquals("abcd", ListFilter.escapeForMySqlRegexReplacement("abcd")); + assertEquals("", ListFilter.escapeForMySqlRegexReplacement("")); + } + + @Test + void test_escapeForMySqlRegexReplacement_doublesApostrophesOnce() { + // Apostrophes only matter for the SQL string-literal layer — REGEXP_REPLACE's + // replacement context doesn't reserve them. Expect a single ' → '' doubling. + assertEquals("a''b", ListFilter.escapeForMySqlRegexReplacement("a'b")); + } + + @Test + void test_escapeForMySqlRegexReplacement_quadruplesBackslashes() { + // One input backslash needs to round-trip to one literal backslash in the + // REGEXP_REPLACE output, so it must be FOUR backslashes in the emitted SQL text: + // SQL text : \\\\ (4 backslashes) + // SQL parser: \\ (2 backslashes — '\\' is the SQL string-literal escape for '\') + // regex eng : \ (1 backslash — '\\' in the regex replacement is a literal '\') + // Without the regex-replacement escape, the regex engine would interpret the lone + // remaining '\' as the start of an escape/backref sequence. + assertEquals("a\\\\\\\\b", ListFilter.escapeForMySqlRegexReplacement("a\\b")); + assertEquals("\\\\\\\\", ListFilter.escapeForMySqlRegexReplacement("\\")); + } + + @Test + void test_escapeForMySqlRegexReplacement_protectsBackreferenceLookalikes() { + // Without the extra regex-replacement layer, "\1" in the input would survive as "\1" + // in the regex replacement and be interpreted as a backreference to capture group 1 + // (REGEXP_REPLACE doesn't have groups when called like updateFqn does, but the + // behaviour is implementation-defined — usually empty-string substitution). After + // the double escape it survives as a literal "\1" in the output. + assertEquals("\\\\\\\\1bar", ListFilter.escapeForMySqlRegexReplacement("\\1bar")); + } + + @Test + void test_escapeForMySqlRegexReplacement_combinesBackslashAndApostrophe() { + // Backslashes get four-x'd, apostrophes double once. + assertEquals("a\\\\\\\\''b", ListFilter.escapeForMySqlRegexReplacement("a\\'b")); + } + @Test void addCondition() { String condition; diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/AsyncServiceTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/AsyncServiceTest.java index 77913e958f11..79c930902c71 100644 --- a/openmetadata-service/src/test/java/org/openmetadata/service/util/AsyncServiceTest.java +++ b/openmetadata-service/src/test/java/org/openmetadata/service/util/AsyncServiceTest.java @@ -17,7 +17,6 @@ import java.util.concurrent.CompletionException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Supplier; @@ -25,7 +24,6 @@ import org.junit.jupiter.api.Test; import org.openmetadata.service.OpenMetadataApplicationConfig; import org.openmetadata.service.OpenMetadataApplicationConfigHolder; -import org.openmetadata.service.jdbi3.HikariCPDataSourceFactory; class AsyncServiceTest { @@ -186,7 +184,6 @@ void testGetInstanceReturnsSingleton() throws Exception { AsyncService second = AsyncService.getInstance(); assertSame(first, second); - assertTrue(first.getMaxConcurrency() >= 4); } @Test @@ -314,44 +311,6 @@ void testExecuteWithRetryInterruptedSleepRestoresInterruptFlag() throws Exceptio Thread.interrupted(); } - @Test - void testResolveMaxConcurrencyUsesConfigBudgetAndCpuFallback() throws Exception { - Method method = AsyncService.class.getDeclaredMethod("resolveMaxConcurrency"); - method.setAccessible(true); - - int cpuBudget = Runtime.getRuntime().availableProcessors() * 2; - setConfigHolderInstance(null); - assertEquals(Integer.valueOf(Math.max(4, cpuBudget)), invoke(method, null)); - - OpenMetadataApplicationConfig config = mock(OpenMetadataApplicationConfig.class); - HikariCPDataSourceFactory dataSourceFactory = mock(HikariCPDataSourceFactory.class); - when(config.getDataSourceFactory()).thenReturn(dataSourceFactory); - when(dataSourceFactory.getMaxSize()).thenReturn(30); - setConfigHolderInstance(config); - - assertEquals(Integer.valueOf(Math.max(4, Math.min(cpuBudget, 10))), invoke(method, null)); - } - - @Test - void testBoundedExecutorLifecycleDelegatesState() throws Exception { - ExecutorService delegate = mock(ExecutorService.class); - when(delegate.isShutdown()).thenReturn(true); - when(delegate.isTerminated()).thenReturn(true); - when(delegate.awaitTermination(5, TimeUnit.SECONDS)).thenReturn(true); - - ExecutorService boundedExecutor = newBoundedExecutorService(delegate); - - assertTrue(boundedExecutor.isShutdown()); - assertTrue(boundedExecutor.isTerminated()); - assertTrue(boundedExecutor.awaitTermination(5, TimeUnit.SECONDS)); - - boundedExecutor.shutdown(); - boundedExecutor.shutdownNow(); - - verify(delegate).shutdown(); - verify(delegate).shutdownNow(); - } - @Test void testShutdownForcesExecutorOnTimeoutAndInterrupt() throws Exception { AsyncService timeoutService = newAsyncService(); @@ -384,16 +343,6 @@ private static AsyncService newAsyncService() throws Exception { return constructor.newInstance(); } - private static ExecutorService newBoundedExecutorService(ExecutorService delegate) - throws Exception { - Class boundedClass = - Class.forName("org.openmetadata.service.util.AsyncService$BoundedExecutorService"); - Constructor constructor = - boundedClass.getDeclaredConstructor(ExecutorService.class, Semaphore.class); - constructor.setAccessible(true); - return (ExecutorService) constructor.newInstance(delegate, new Semaphore(1)); - } - private static void replaceExecutor(AsyncService service, ExecutorService executor) throws Exception { ExecutorService originalExecutor = service.getExecutorService();