fix(ingestion): Prevent duplicate [CLI] source creation during executor-managed ingestion (#16663)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: treff7es <treff7es@gmail.com>

Author: 3 people

metadata-ingestion/src/datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -28,7 +28,9 @@
     StructuredExecutionReportClass,
     _Aspect,
 )
+from datahub.metadata.urns import DataHubExecutionRequestUrn
 from datahub.utilities.logging_manager import get_log_buffer
+from datahub.utilities.urns.error import InvalidUrnError
 from datahub.utilities.urns.urn import Urn
 
 logger = logging.getLogger(__name__)
@@ -118,35 +120,63 @@ def __init__(self, sink: Sink, report_recipe: bool, ctx: PipelineContext) -> Non
         ingestion_source_key = self.generate_unique_key(ctx.pipeline_config)
         self.entity_name: str = self.generate_entity_name(ingestion_source_key)
 
+        # If run_id is an execution request URN, the executor owns the source/request lifecycle.
+        try:
+            parsed = Urn.from_string(ctx.run_id)
+            self._is_running_under_executor = (
+                parsed.entity_type == DataHubExecutionRequestUrn.ENTITY_TYPE
+            )
+        except InvalidUrnError:
+            self._is_running_under_executor = False
+        except Exception:
+            logger.warning(
+                f"Unexpected error parsing run_id={ctx.run_id!r} as URN; "
+                "assuming standalone CLI context.",
+                exc_info=True,
+            )
+            self._is_running_under_executor = False
+
+        if self._is_running_under_executor:
+            logger.debug(f"Executor-managed run detected (run_id={ctx.run_id}).")
+
         self.ingestion_source_urn: Urn = Urn(
             entity_type="dataHubIngestionSource",
             entity_id=["cli-" + datahub_guid(ingestion_source_key)],
         )
         logger.debug(f"Ingestion source urn = {self.ingestion_source_urn}")
-        self.execution_request_input_urn: Urn = Urn(
-            entity_type="dataHubExecutionRequest", entity_id=[ctx.run_id]
-        )
+        # Use typed URN only in the executor path (run_id already validated as such).
+        # For standalone CLI runs, run_id is a plain string; passing a foreign URN type
+        # to DataHubExecutionRequestUrn would raise InvalidUrnError.
+        if self._is_running_under_executor:
+            self.execution_request_input_urn: Urn = DataHubExecutionRequestUrn(
+                ctx.run_id
+            )
+        else:
+            self.execution_request_input_urn = Urn(
+                entity_type="dataHubExecutionRequest", entity_id=[ctx.run_id]
+            )
         self.start_time_ms: int = self.get_cur_time_in_ms()
 
-        # Construct the dataHubIngestionSourceInfo aspect
-        source_info_aspect = DataHubIngestionSourceInfoClass(
-            name=self.entity_name,
-            type=ctx.pipeline_config.source.type,
-            platform=make_data_platform_urn(
-                getattr(ctx.pipeline_config.source, "platform", "unknown")
-            ),
-            config=DataHubIngestionSourceConfigClass(
-                recipe=self._get_recipe_to_report(ctx),
-                version=nice_version_name(),
-                executorId=self._EXECUTOR_ID,
-            ),
-        )
+        if not self._is_running_under_executor:
+            # Construct the dataHubIngestionSourceInfo aspect
+            source_info_aspect = DataHubIngestionSourceInfoClass(
+                name=self.entity_name,
+                type=ctx.pipeline_config.source.type,
+                platform=make_data_platform_urn(
+                    getattr(ctx.pipeline_config.source, "platform", "unknown")
+                ),
+                config=DataHubIngestionSourceConfigClass(
+                    recipe=self._get_recipe_to_report(ctx),
+                    version=nice_version_name(),
+                    executorId=self._EXECUTOR_ID,
+                ),
+            )
 
-        # Emit the dataHubIngestionSourceInfo aspect
-        self._emit_aspect(
-            entity_urn=self.ingestion_source_urn,
-            aspect_value=source_info_aspect,
-        )
+            # Emit the dataHubIngestionSourceInfo aspect
+            self._emit_aspect(
+                entity_urn=self.ingestion_source_urn,
+                aspect_value=source_info_aspect,
+            )
 
     @staticmethod
     def _convert_sets_to_lists(obj: Any) -> Any:
@@ -214,6 +244,10 @@ def _emit_aspect(
 
     def on_start(self, ctx: PipelineContext) -> None:
         assert ctx.pipeline_config is not None
+
+        if self._is_running_under_executor:
+            return
+
         # Construct the dataHubExecutionRequestInput aspect
         execution_input_aspect = ExecutionRequestInputClass(
             task=self._INGESTION_TASK_NAME,

metadata-ingestion/tests/unit/reporting/test_datahub_ingestion_reporter.py

@@ -1,12 +1,19 @@
 from typing import Any, Dict, List, Set, Tuple, Union
+from unittest.mock import MagicMock
 
 import pytest
 
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.sink import Sink
 from datahub.ingestion.reporting.datahub_ingestion_run_summary_provider import (
     DatahubIngestionRunSummaryProvider,
     DatahubIngestionRunSummaryProviderConfig,
 )
 from datahub.ingestion.run.pipeline_config import PipelineConfig
+from datahub.metadata.schema_classes import (
+    DataHubIngestionSourceInfoClass,
+    ExecutionRequestInputClass,
+)
 
 
 @pytest.mark.parametrize(
@@ -185,3 +192,62 @@ def test_json_serializable() -> None:
     except TypeError:
         serializable = False
     assert serializable
+
+
+def _make_provider(run_id: str, report_recipe: bool = True) -> tuple:
+    pipeline_config = PipelineConfig.from_dict(
+        {"source": {"type": "bigquery"}, "sink": {"type": "console"}}
+    )
+    ctx = PipelineContext(run_id=run_id, pipeline_config=pipeline_config)
+    mock_sink = MagicMock(spec=Sink)
+    provider = DatahubIngestionRunSummaryProvider(
+        sink=mock_sink, report_recipe=report_recipe, ctx=ctx
+    )
+    return provider, ctx, mock_sink
+
+
+def test_executor_detection_with_execution_request_urn() -> None:
+    """Test that CLI source creation is skipped when run_id is an execution request URN"""
+    _provider, _ctx, mock_sink = _make_provider("urn:li:dataHubExecutionRequest:12345")
+
+    mock_sink.write_record_async.assert_not_called()
+
+
+def test_cli_source_creation_with_normal_run_id() -> None:
+    """Test that CLI source is created when run_id is NOT an execution request URN"""
+    provider, _ctx, mock_sink = _make_provider("normal-cli-run-id-12345")
+
+    assert mock_sink.write_record_async.call_count == 1
+    record_envelope = mock_sink.write_record_async.call_args[0][0]
+    assert record_envelope.record.entityUrn == str(provider.ingestion_source_urn)
+    assert isinstance(record_envelope.record.aspect, DataHubIngestionSourceInfoClass)
+
+
+def test_on_start_skips_execution_request_input_under_executor() -> None:
+    """Test that on_start skips ExecutionRequestInput creation when under executor"""
+    provider, ctx, mock_sink = _make_provider("urn:li:dataHubExecutionRequest:12345")
+    mock_sink.reset_mock()  # Reset after __init__ (which emits nothing in executor path)
+
+    provider.on_start(ctx)
+
+    mock_sink.write_record_async.assert_not_called()
+
+
+def test_on_start_creates_execution_request_input_for_cli() -> None:
+    """Test that on_start creates ExecutionRequestInput for standalone CLI runs"""
+    provider, ctx, mock_sink = _make_provider("normal-cli-run-id")
+    mock_sink.reset_mock()  # Reset after __init__
+
+    provider.on_start(ctx)
+
+    assert mock_sink.write_record_async.call_count == 1
+    record_envelope = mock_sink.write_record_async.call_args[0][0]
+    assert isinstance(record_envelope.record.aspect, ExecutionRequestInputClass)
+
+
+def test_non_execution_request_urn_as_run_id_does_not_crash() -> None:
+    """A valid URN of a different entity type must not crash __init__."""
+    provider, _ctx, _mock_sink = _make_provider(
+        "urn:li:ingestionSource:some-source-id", report_recipe=False
+    )
+    assert not provider._is_running_under_executor