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)