From 1e56c76c0efc62cbe45d3ebfbd26a139f1126d20 Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull <peremiquelbrull@gmail.com>
Date: Fri, 20 Sep 2024 10:37:41 +0200
Subject: [PATCH] FIX #17896 - Python lineage SDK to work with Uuid & FQN
 models (#17928)

---
 .../ingestion/ometa/mixins/lineage_mixin.py   | 14 +++++++----
 .../ometa/test_ometa_lineage_api.py           | 24 +++++++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/ometa/mixins/lineage_mixin.py b/ingestion/src/metadata/ingestion/ometa/mixins/lineage_mixin.py
index ad342276ac7..45769f5935c 100644
--- a/ingestion/src/metadata/ingestion/ometa/mixins/lineage_mixin.py
+++ b/ingestion/src/metadata/ingestion/ometa/mixins/lineage_mixin.py
@@ -22,13 +22,14 @@ from pydantic import BaseModel
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
+from metadata.generated.schema.type.basic import FullyQualifiedEntityName, Uuid
 from metadata.generated.schema.type.entityLineage import ColumnLineage, EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.lineage.models import ConnectionTypeDialectMapper
 from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT
 from metadata.ingestion.models.patch_request import build_patch
 from metadata.ingestion.ometa.client import REST, APIError
-from metadata.ingestion.ometa.utils import get_entity_type, quote
+from metadata.ingestion.ometa.utils import get_entity_type, model_str, quote
 from metadata.utils.logger import ometa_logger
 from metadata.utils.lru_cache import LRU_CACHE_SIZE, LRUCache
 
@@ -260,7 +261,7 @@ class OMetaLineageMixin(Generic[T]):
     def get_lineage_by_id(
         self,
         entity: Union[Type[T], str],
-        entity_id: str,
+        entity_id: Union[str, Uuid],
         up_depth: int = 1,
         down_depth: int = 1,
     ) -> Optional[Dict[str, Any]]:
@@ -272,13 +273,16 @@ class OMetaLineageMixin(Generic[T]):
         :param down_depth: Downstream depth of lineage (default=1, min=0, max=3)
         """
         return self._get_lineage(
-            entity=entity, path=entity_id, up_depth=up_depth, down_depth=down_depth
+            entity=entity,
+            path=model_str(entity_id),
+            up_depth=up_depth,
+            down_depth=down_depth,
         )
 
     def get_lineage_by_name(
         self,
         entity: Union[Type[T], str],
-        fqn: str,
+        fqn: Union[str, FullyQualifiedEntityName],
         up_depth: int = 1,
         down_depth: int = 1,
     ) -> Optional[Dict[str, Any]]:
@@ -291,7 +295,7 @@ class OMetaLineageMixin(Generic[T]):
         """
         return self._get_lineage(
             entity=entity,
-            path=f"name/{quote(fqn)}",
+            path=f"name/{quote(model_str(fqn))}",
             up_depth=up_depth,
             down_depth=down_depth,
         )
diff --git a/ingestion/tests/integration/ometa/test_ometa_lineage_api.py b/ingestion/tests/integration/ometa/test_ometa_lineage_api.py
index b661d6e2bad..f083a50f6cc 100644
--- a/ingestion/tests/integration/ometa/test_ometa_lineage_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_lineage_api.py
@@ -288,6 +288,30 @@ class OMetaLineageTest(TestCase):
             len(res["downstreamEdges"][0]["lineageDetails"]["columnsLineage"]), 2
         )
 
+        # We can get lineage by ID
+        lineage_id = self.metadata.get_lineage_by_id(
+            entity=Table, entity_id=self.table2_entity.id.root
+        )
+        assert lineage_id["entity"]["id"] == str(self.table2_entity.id.root)
+
+        # Same thing works if we pass directly the Uuid
+        lineage_uuid = self.metadata.get_lineage_by_id(
+            entity=Table, entity_id=self.table2_entity.id
+        )
+        assert lineage_uuid["entity"]["id"] == str(self.table2_entity.id.root)
+
+        # We can also get lineage by name
+        lineage_str = self.metadata.get_lineage_by_name(
+            entity=Table, fqn=self.table2_entity.fullyQualifiedName.root
+        )
+        assert lineage_str["entity"]["id"] == str(self.table2_entity.id.root)
+
+        # Or passing the FQN
+        lineage_fqn = self.metadata.get_lineage_by_name(
+            entity=Table, fqn=self.table2_entity.fullyQualifiedName
+        )
+        assert lineage_fqn["entity"]["id"] == str(self.table2_entity.id.root)
+
     def test_delete_by_source(self):
         """
         Test case for deleting lineage by source.