From 309abf0c94cb45fefe0971011f6a919557b09bf7 Mon Sep 17 00:00:00 2001
From: Hyejin Yoon <0327jane@gmail.com>
Date: Thu, 19 Dec 2024 22:16:08 +0900
Subject: [PATCH] update mcp.py

---
 metadata-ingestion/src/datahub/emitter/mcp.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/metadata-ingestion/src/datahub/emitter/mcp.py b/metadata-ingestion/src/datahub/emitter/mcp.py
index c6fcfad2e0abaa..60c4d6b4d9d457 100644
--- a/metadata-ingestion/src/datahub/emitter/mcp.py
+++ b/metadata-ingestion/src/datahub/emitter/mcp.py
@@ -14,7 +14,7 @@
     SystemMetadataClass,
     _Aspect,
 )
-from datahub.utilities.urns.urn import guess_entity_type
+from datahub.utilities.urns.urn import Urn, guess_entity_type
 
 if TYPE_CHECKING:
     from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -63,7 +63,7 @@ class MetadataChangeProposalWrapper:
 
     entityType: str = _ENTITY_TYPE_UNSET
     changeType: Union[str, ChangeTypeClass] = ChangeTypeClass.UPSERT
-    entityUrn: Union[None, str] = None
+    entityUrn: Union[None, str,] = None
     entityKeyAspect: Union[None, _Aspect] = None
     auditHeader: Union[None, KafkaAuditHeaderClass] = None
     aspectName: Union[None, str] = None
@@ -71,7 +71,11 @@ class MetadataChangeProposalWrapper:
     systemMetadata: Union[None, SystemMetadataClass] = None
 
     def __post_init__(self) -> None:
-        if self.entityUrn and self.entityType == _ENTITY_TYPE_UNSET:
+        if isinstance(self.entityUrn, Urn):
+            if self.entityType == _ENTITY_TYPE_UNSET:
+                self.entityType = self.entityUrn.entity_type
+            self.entityUrn = str(self.entityUrn)
+        elif self.entityUrn and self.entityType == _ENTITY_TYPE_UNSET:
             self.entityType = guess_entity_type(self.entityUrn)
         elif self.entityUrn and self.entityType:
             guessed_entity_type = guess_entity_type(self.entityUrn).lower()
@@ -104,7 +108,8 @@ def construct_many(
     ) -> List["MetadataChangeProposalWrapper"]:
         return [cls(entityUrn=entityUrn, aspect=aspect) for aspect in aspects if aspect]
 
-    def _make_mcp_without_aspects(self) -> MetadataChangeProposalClass:
+    def _make_mcp_without_aspects(self) -> MetadataChangeProposalClass: 
+        assert self.entityUrn is None or isinstance(self.entityUrn, str)
         return MetadataChangeProposalClass(
             entityType=self.entityType,
             entityUrn=self.entityUrn,