fix(ingest): pass platform correctly to browse path v2 helper (#8244)

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

@@ -23,7 +23,7 @@ from pydantic import BaseModel
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.source_common import PlatformInstanceConfigMixin
-from datahub.emitter.mcp_builder import PlatformKey, mcps_from_mce
+from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope, WorkUnit
 from datahub.ingestion.api.report import Report
@@ -249,19 +249,14 @@ class Source(Closeable, metaclass=ABCMeta):
             env and env.lower(),
         ]
 
-        platform_key: Optional[PlatformKey] = None
-        if (
-            platform
-            and isinstance(config, PlatformInstanceConfigMixin)
-            and config.platform_instance
-        ):
-            platform_key = PlatformKey(
-                platform=platform, instance=config.platform_instance
-            )
+        platform_instance: Optional[str] = None
+        if isinstance(config, PlatformInstanceConfigMixin) and config.platform_instance:
+            platform_instance = platform_instance
 
         return partial(
             auto_browse_path_v2,
-            platform_key=platform_key,
+            platform=platform,
+            platform_instance=platform_instance,
             drop_dirs=[s for s in browse_path_drop_dirs if s is not None],
             dry_run=dry_run,
         )

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

@@ -15,7 +15,6 @@ from typing import (
 
 from datahub.emitter.mce_builder import make_dataplatform_instance_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.mcp_builder import PlatformKey
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
     BrowsePathEntryClass,
@@ -173,7 +172,8 @@ def auto_browse_path_v2(
     *,
     dry_run: bool = False,
     drop_dirs: Sequence[str] = (),
-    platform_key: Optional[PlatformKey] = None,
+    platform: Optional[str] = None,
+    platform_instance: Optional[str] = None,
 ) -> Iterable[MetadataWorkUnit]:
     """Generate BrowsePathsV2 from Container and BrowsePaths aspects.
 
@@ -248,7 +248,9 @@ def auto_browse_path_v2(
                 yield MetadataChangeProposalWrapper(
                     entityUrn=urn,
                     aspect=BrowsePathsV2Class(
-                        path=_prepend_platform_instance(path, platform_key)
+                        path=_prepend_platform_instance(
+                            path, platform, platform_instance
+                        )
                     ),
                 ).as_workunit()
         elif urn not in emitted_urns and guess_entity_type(urn) == "container":
@@ -258,16 +260,14 @@ def auto_browse_path_v2(
                 yield MetadataChangeProposalWrapper(
                     entityUrn=urn,
                     aspect=BrowsePathsV2Class(
-                        path=_prepend_platform_instance([], platform_key)
+                        path=_prepend_platform_instance([], platform, platform_instance)
                     ),
                 ).as_workunit()
 
     if num_out_of_batch or num_out_of_order:
         properties = {
-            "platform": platform_key.platform if platform_key else None,
-            "has_platform_instance": bool(platform_key.instance)
-            if platform_key
-            else False,
+            "platform": platform,
+            "has_platform_instance": bool(platform_instance),
             "num_out_of_batch": num_out_of_batch,
             "num_out_of_order": num_out_of_order,
         }
@@ -293,12 +293,12 @@ def _batch_workunits_by_urn(
 
 
 def _prepend_platform_instance(
-    entries: List[BrowsePathEntryClass], platform_key: Optional[PlatformKey]
+    entries: List[BrowsePathEntryClass],
+    platform: Optional[str],
+    platform_instance: Optional[str],
 ) -> List[BrowsePathEntryClass]:
-    if platform_key and platform_key.instance:
-        urn = make_dataplatform_instance_urn(
-            platform_key.platform, platform_key.instance
-        )
+    if platform and platform_instance:
+        urn = make_dataplatform_instance_urn(platform, platform_instance)
         return [BrowsePathEntryClass(id=urn, urn=urn)] + entries
 
     return entries

metadata-ingestion/tests/unit/test_source_helpers.py

@@ -8,7 +8,6 @@ from datahub.emitter.mce_builder import (
     make_dataset_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.mcp_builder import PlatformKey
 from datahub.ingestion.api.source_helpers import (
     auto_browse_path_v2,
     auto_status_aspect,
@@ -298,7 +297,6 @@ def test_auto_browse_path_v2_container_over_legacy_browse_path(telemetry_ping_mo
 def test_auto_browse_path_v2_with_platform_instsance(telemetry_ping_mock):
     platform = "my_platform"
     platform_instance = "my_instance"
-    platform_key = PlatformKey(platform=platform, instance=platform_instance)
     platform_instance_urn = make_dataplatform_instance_urn(platform, platform_instance)
     platform_instance_entry = models.BrowsePathEntryClass(
         platform_instance_urn, platform_instance_urn
@@ -310,7 +308,8 @@ def test_auto_browse_path_v2_with_platform_instsance(telemetry_ping_mock):
     new_wus = list(
         auto_browse_path_v2(
             wus,
-            platform_key=platform_key,
+            platform=platform,
+            platform_instance=platform_instance,
         )
     )
     assert telemetry_ping_mock.call_count == 0