fix(ingest): only add to samples where platform match (#14722)

Author: anshbansal

metadata-ingestion/scripts/avro_codegen.py

@@ -469,6 +469,10 @@ def create_from_ids(cls, data_flow_urn: str, job_id: str) -> "DataJobUrn":
 def get_data_flow_urn(self) -> "DataFlowUrn":
     return DataFlowUrn.from_string(self.flow)
 
+@property
+def orchestrator(self) -> str:
+    return self.get_data_flow_urn().orchestrator
+
 @deprecated(reason="Use .job_id instead")
 def get_job_id(self) -> str:
     return self.job_id

metadata-ingestion/src/datahub/ingestion/api/report.py

@@ -29,6 +29,7 @@
 )
 from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.lossy_collections import LossyList
+from datahub.utilities.urns.urn import guess_platform_name
 
 logger = logging.getLogger(__name__)
 LogLevel = Literal["ERROR", "WARNING", "INFO", "DEBUG"]
@@ -41,6 +42,15 @@ def as_obj(self) -> dict: ...
 
 @dataclass
 class Report(SupportsAsObj):
+    def __post_init__(self) -> None:
+        self.platform: Optional[str] = None
+
+    def set_platform(self, platform: str) -> None:
+        self.platform = platform
+
+    def get_platform(self) -> Optional[str]:
+        return self.platform
+
     @staticmethod
     def to_str(some_val: Any) -> str:
         if isinstance(some_val, Enum):
@@ -213,6 +223,7 @@ class ExamplesReport(Report, Closeable):
     _lineage_aspects_seen: Set[str] = field(default_factory=set)
 
     def __post_init__(self) -> None:
+        super().__post_init__()
         self._file_based_dict = FileBackedDict(
             tablename="urn_aspects",
             extra_columns={
@@ -347,6 +358,9 @@ def _update_file_based_dict(
         aspectName: str,
         mcp: Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper],
     ) -> None:
+        platform_name = guess_platform_name(urn)
+        if platform_name != self.get_platform():
+            return
         if is_lineage_aspect(entityType, aspectName):
             self._lineage_aspects_seen.add(aspectName)
         has_fine_grained_lineage = self._has_fine_grained_lineage(mcp)

metadata-ingestion/src/datahub/ingestion/api/source.py

@@ -531,9 +531,9 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
             auto_status_aspect,
             auto_materialize_referenced_tags_terms,
             partial(
-                auto_fix_duplicate_schema_field_paths, platform=self._infer_platform()
+                auto_fix_duplicate_schema_field_paths, platform=self.infer_platform()
             ),
-            partial(auto_fix_empty_field_paths, platform=self._infer_platform()),
+            partial(auto_fix_empty_field_paths, platform=self.infer_platform()),
             browse_path_processor,
             partial(auto_workunit_reporter, self.get_report()),
             auto_patch_last_modified,
@@ -583,7 +583,7 @@ def get_report(self) -> SourceReport:
     def close(self) -> None:
         self.get_report().close()
 
-    def _infer_platform(self) -> Optional[str]:
+    def infer_platform(self) -> Optional[str]:
         config = self.get_config()
         platform = (
             getattr(config, "platform_name", None)
@@ -598,7 +598,7 @@ def _infer_platform(self) -> Optional[str]:
     def _get_browse_path_processor(self, dry_run: bool) -> MetadataWorkUnitProcessor:
         config = self.get_config()
 
-        platform = self._infer_platform()
+        platform = self.infer_platform()
         env = getattr(config, "env", None)
         browse_path_drop_dirs = [
             platform,

metadata-ingestion/src/datahub/ingestion/run/pipeline.py

@@ -440,7 +440,19 @@ def _time_to_print(self) -> bool:
             return True
         return False
 
+    def _set_platform(self) -> None:
+        platform = self.source.infer_platform()
+        if platform:
+            self.source.get_report().set_platform(platform)
+        else:
+            self.source.get_report().warning(
+                message="Platform not found",
+                title="Platform not found",
+                context="Platform not found",
+            )
+
     def run(self) -> None:
+        self._set_platform()
         self._warn_old_cli_version()
         with self.exit_stack, self.inner_exit_stack:
             if self.config.flags.generate_memory_profiles:

metadata-ingestion/src/datahub/ingestion/source/iceberg/iceberg.py

@@ -200,9 +200,9 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
             auto_lowercase_dataset_urns,
             auto_materialize_referenced_tags_terms,
             partial(
-                auto_fix_duplicate_schema_field_paths, platform=self._infer_platform()
+                auto_fix_duplicate_schema_field_paths, platform=self.infer_platform()
             ),
-            partial(auto_fix_empty_field_paths, platform=self._infer_platform()),
+            partial(auto_fix_empty_field_paths, platform=self.infer_platform()),
             partial(auto_workunit_reporter, self.get_report()),
             auto_patch_last_modified,
             EnsureAspectSizeProcessor(self.get_report()).ensure_aspect_size,

metadata-ingestion/src/datahub/utilities/urns/urn.py

@@ -1,8 +1,47 @@
-from datahub.metadata.urns import Urn
+from typing import Optional
 
-__all__ = ["Urn", "guess_entity_type"]
+from datahub.metadata.urns import (
+    DataPlatformUrn,
+    Urn,
+)
+
+__all__ = ["Urn", "guess_entity_type", "guess_platform_name"]
 
 
 def guess_entity_type(urn: str) -> str:
     assert urn.startswith("urn:li:"), "urns must start with urn:li:"
     return urn.split(":")[2]
+
+
+def guess_platform_name(urn: str) -> Optional[str]:
+    """Extract platform from URN using a mapping dictionary."""
+    urn_obj = Urn.from_string(urn)
+
+    try:
+        platform = None
+        try:
+            platform = urn_obj.platform  # type: ignore[attr-defined]
+            platform_name = DataPlatformUrn.from_string(
+                platform
+            ).get_entity_id_as_string()
+            return platform_name
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.orchestrator  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.dashboard_tool  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+        try:
+            return urn_obj.ml_model_tool  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+
+        if platform is None:
+            return None
+    except AttributeError:
+        pass
+    return None

metadata-ingestion/tests/unit/api/test_pipeline.py

@@ -9,6 +9,7 @@
 from datahub.configuration.common import DynamicTypedConfig
 from datahub.ingestion.api.committable import CommitPolicy, Committable
 from datahub.ingestion.api.common import RecordEnvelope
+from datahub.ingestion.api.decorators import platform_name
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -514,6 +515,7 @@ def transform(
             yield record_envelope
 
 
+@platform_name("fake")
 class FakeSource(Source):
     def __init__(self, ctx: PipelineContext):
         super().__init__(ctx)
@@ -537,6 +539,7 @@ def close(self):
         pass
 
 
+@platform_name("fake")
 class FakeSourceWithWarnings(FakeSource):
     def __init__(self, ctx: PipelineContext):
         super().__init__(ctx)
@@ -546,6 +549,7 @@ def get_report(self) -> SourceReport:
         return self.source_report
 
 
+@platform_name("fake")
 class FakeSourceWithFailures(FakeSource):
     def __init__(self, ctx: PipelineContext):
         super().__init__(ctx)

metadata-ingestion/tests/unit/datahub/utilities/__init__.py

@@ -0,0 +1,26 @@
+from datahub.utilities.urns.urn import guess_platform_name
+
+
+def test_guess_platform_name():
+    assert guess_platform_name("urn:li:corpuser:jdoe") is None
+    assert (
+        guess_platform_name(
+            "urn:li:dataset:(urn:li:dataPlatform:kafka,SampleKafkaDataset,PROD)"
+        )
+        == "kafka"
+    )
+    assert guess_platform_name("urn:li:dataFlow:(airflow,dag_abc,PROD)") == "airflow"
+    assert (
+        guess_platform_name(
+            "urn:li:dataJob:(urn:li:dataFlow:(airflow,dag_abc,PROD),task_123)"
+        )
+        == "airflow"
+    )
+    assert guess_platform_name("urn:li:chart:(looker,baz1)") == "looker"
+    assert guess_platform_name("urn:li:dashboard:(looker,baz)") == "looker"
+    assert (
+        guess_platform_name(
+            "urn:li:mlModel:(urn:li:dataPlatform:science,scienceModel,PROD)"
+        )
+        == "science"
+    )

metadata-ingestion/tests/unit/datahub/utilities/test_urn.py

@@ -2,6 +2,7 @@
 
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import platform_name
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
@@ -24,13 +25,14 @@
 def _get_urn(table_name: str = "fooIndex") -> str:
     return str(
         DatasetUrn.create_from_ids(
-            platform_id="elasticsearch",
+            platform_id="fake",
             table_name=table_name,
             env="PROD",
         )
     )
 
 
+@platform_name("fake")
 class FakeSource(Source):
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         return [
@@ -46,6 +48,7 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
     def __init__(self, ctx: PipelineContext):
         super().__init__(ctx)
         self.source_report = SourceReport()
+        self.source_report.set_platform("fake")
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "FakeSource":