From 185ef5e98f09ea54b477ccc1daabf1a4761812a1 Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull
Date: Mon, 6 Dec 2021 08:40:53 +0100
Subject: [PATCH] Get id as str or uuid (#1562)
---
.../src/metadata/ingestion/ometa/ometa_api.py | 32 +++++++++++++++----
.../integration/ometa/test_ometa_chart_api.py | 2 +-
.../ometa/test_ometa_dashboard_api.py | 4 +--
.../ometa/test_ometa_database_api.py | 4 +--
.../integration/ometa/test_ometa_model_api.py | 4 +--
.../ometa/test_ometa_pipeline_api.py | 4 +--
.../integration/ometa/test_ometa_table_api.py | 4 +--
.../integration/ometa/test_ometa_topic_api.py | 2 +-
8 files changed, 33 insertions(+), 23 deletions(-)
diff --git a/ingestion/src/metadata/ingestion/ometa/ometa_api.py b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
index 4aa6a29072d..8b2a21465d1 100644
--- a/ingestion/src/metadata/ingestion/ometa/ometa_api.py
+++ b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
@@ -34,6 +34,7 @@ from metadata.generated.schema.entity.services.storageService import StorageServ
from metadata.generated.schema.entity.tags.tagCategory import Tag
from metadata.generated.schema.entity.teams.team import Team
from metadata.generated.schema.entity.teams.user import User
+from metadata.generated.schema.type import basic
from metadata.generated.schema.type.entityHistory import EntityVersionHistory
from metadata.ingestion.ometa.auth_provider import AuthenticationProvider
from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
@@ -336,6 +337,20 @@ class OpenMetadata(OMetaLineageMixin, OMetaTableMixin, Generic[T, C]):
resp = method(self.get_suffix(entity), data=data.json())
return entity_class(**resp)
+ @staticmethod
+ def uuid_to_str(entity_id: Union[str, basic.Uuid]) -> str:
+ """
+ Given an entity_id, that can be a str or our pydantic
+ definition of Uuid, return a proper str to build
+ the endpoint path
+ :param entity_id: Entity ID to onvert to string
+ :return: str for the ID
+ """
+ if isinstance(entity_id, basic.Uuid):
+ return str(entity_id.__root__)
+
+ return entity_id
+
def get_by_name(
self, entity: Type[T], fqdn: str, fields: Optional[List[str]] = None
) -> Optional[T]:
@@ -346,13 +361,16 @@ class OpenMetadata(OMetaLineageMixin, OMetaTableMixin, Generic[T, C]):
return self._get(entity=entity, path=f"name/{fqdn}", fields=fields)
def get_by_id(
- self, entity: Type[T], entity_id: str, fields: Optional[List[str]] = None
+ self,
+ entity: Type[T],
+ entity_id: Union[str, basic.Uuid],
+ fields: Optional[List[str]] = None,
) -> Optional[T]:
"""
Return entity by ID or None
"""
- return self._get(entity=entity, path=entity_id, fields=fields)
+ return self._get(entity=entity, path=self.uuid_to_str(entity_id), fields=fields)
def _get(
self, entity: Type[T], path: str, fields: Optional[List[str]] = None
@@ -399,13 +417,15 @@ class OpenMetadata(OMetaLineageMixin, OMetaTableMixin, Generic[T, C]):
after = resp["paging"]["after"] if "after" in resp["paging"] else None
return EntityList(entities=entities, total=total, after=after)
- def list_versions(self, entity_id: str, entity: Type[T]) -> EntityVersionHistory:
+ def list_versions(
+ self, entity_id: Union[str, basic.Uuid], entity: Type[T]
+ ) -> EntityVersionHistory:
"""
Helps us paginate over the collection
"""
suffix = self.get_suffix(entity)
- path = f"/{entity_id}/versions"
+ path = f"/{self.uuid_to_str(entity_id)}/versions"
resp = self.client.get(f"{suffix}{path}")
if self._use_raw_data:
@@ -424,8 +444,8 @@ class OpenMetadata(OMetaLineageMixin, OMetaTableMixin, Generic[T, C]):
else:
return [entity(**p) for p in resp["data"]]
- def delete(self, entity: Type[T], entity_id: str) -> None:
- self.client.delete(f"{self.get_suffix(entity)}/{entity_id}")
+ def delete(self, entity: Type[T], entity_id: Union[str, basic.Uuid]) -> None:
+ self.client.delete(f"{self.get_suffix(entity)}/{self.uuid_to_str(entity_id)}")
def compute_percentile(self, entity: Union[Type[T], str], date: str) -> None:
"""
diff --git a/ingestion/tests/integration/ometa/test_ometa_chart_api.py b/ingestion/tests/integration/ometa/test_ometa_chart_api.py
index a19e955d0b3..ae0f7f6d370 100644
--- a/ingestion/tests/integration/ometa/test_ometa_chart_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_chart_api.py
@@ -147,7 +147,7 @@ class OMetaChartTest(TestCase):
entity=Chart, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(entity=Chart, entity_id=str(res_name.id.__root__))
+ res = self.metadata.get_by_id(entity=Chart, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)
diff --git a/ingestion/tests/integration/ometa/test_ometa_dashboard_api.py b/ingestion/tests/integration/ometa/test_ometa_dashboard_api.py
index 16049d75b4e..376aa6184e8 100644
--- a/ingestion/tests/integration/ometa/test_ometa_dashboard_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_dashboard_api.py
@@ -149,9 +149,7 @@ class OMetaDashboardTest(TestCase):
entity=Dashboard, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(
- entity=Dashboard, entity_id=str(res_name.id.__root__)
- )
+ res = self.metadata.get_by_id(entity=Dashboard, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)
diff --git a/ingestion/tests/integration/ometa/test_ometa_database_api.py b/ingestion/tests/integration/ometa/test_ometa_database_api.py
index f6f058b0435..74337f1dfc6 100644
--- a/ingestion/tests/integration/ometa/test_ometa_database_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_database_api.py
@@ -149,9 +149,7 @@ class OMetaDatabaseTest(TestCase):
entity=Database, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(
- entity=Database, entity_id=str(res_name.id.__root__)
- )
+ res = self.metadata.get_by_id(entity=Database, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)
diff --git a/ingestion/tests/integration/ometa/test_ometa_model_api.py b/ingestion/tests/integration/ometa/test_ometa_model_api.py
index ae32affaf47..9d327c3dabb 100644
--- a/ingestion/tests/integration/ometa/test_ometa_model_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_model_api.py
@@ -121,9 +121,7 @@ class OMetaModelTest(TestCase):
entity=MlModel, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(
- entity=MlModel, entity_id=str(res_name.id.__root__)
- )
+ res = self.metadata.get_by_id(entity=MlModel, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)
diff --git a/ingestion/tests/integration/ometa/test_ometa_pipeline_api.py b/ingestion/tests/integration/ometa/test_ometa_pipeline_api.py
index d105fcca41a..8b04dd807bc 100644
--- a/ingestion/tests/integration/ometa/test_ometa_pipeline_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_pipeline_api.py
@@ -149,9 +149,7 @@ class OMetaPipelineTest(TestCase):
entity=Pipeline, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(
- entity=Pipeline, entity_id=str(res_name.id.__root__)
- )
+ res = self.metadata.get_by_id(entity=Pipeline, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)
diff --git a/ingestion/tests/integration/ometa/test_ometa_table_api.py b/ingestion/tests/integration/ometa/test_ometa_table_api.py
index f203a627b28..bd1223f05b9 100644
--- a/ingestion/tests/integration/ometa/test_ometa_table_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_table_api.py
@@ -212,9 +212,7 @@ class OMetaTableTest(TestCase):
entity=Table, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res_id = self.metadata.get_by_id(
- entity=Table, entity_id=str(res_name.id.__root__)
- )
+ res_id = self.metadata.get_by_id(entity=Table, entity_id=res_name.id)
# Delete
self.metadata.delete(entity=Table, entity_id=str(res_id.id.__root__))
diff --git a/ingestion/tests/integration/ometa/test_ometa_topic_api.py b/ingestion/tests/integration/ometa/test_ometa_topic_api.py
index 2ab25e7592c..4f12a73c5bf 100644
--- a/ingestion/tests/integration/ometa/test_ometa_topic_api.py
+++ b/ingestion/tests/integration/ometa/test_ometa_topic_api.py
@@ -149,7 +149,7 @@ class OMetaTopicTest(TestCase):
entity=Topic, fqdn=self.entity.fullyQualifiedName
)
# Then fetch by ID
- res = self.metadata.get_by_id(entity=Topic, entity_id=str(res_name.id.__root__))
+ res = self.metadata.get_by_id(entity=Topic, entity_id=res_name.id)
self.assertEqual(res_name.id, res.id)