diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py index ef56e90c32..996f2a0f0d 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/proxy.py @@ -134,6 +134,13 @@ class UnityCatalogApiProxy: self._unity_catalog_api.list_metastores() return True + def assigned_metastore(self) -> Optional[Metastore]: + response: dict = self._unity_catalog_api.get_metastore_summary() + if response.get("metastore_id") is None: + logger.info("Not found assigned metastore") + return None + return self._create_metastore(response) + def metastores(self) -> Iterable[Metastore]: response: dict = self._unity_catalog_api.list_metastores() if response.get("metastores") is None: diff --git a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py index c334da6d15..afd87e2009 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/unity/source.py +++ b/metadata-ingestion/src/datahub/ingestion/source/unity/source.py @@ -1,6 +1,6 @@ import logging import re -from typing import Iterable, List, Optional +from typing import Dict, Iterable, List, Optional from datahub.emitter.mce_builder import ( make_data_platform_urn, @@ -172,7 +172,14 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource): yield from self.process_metastores() def process_metastores(self) -> Iterable[MetadataWorkUnit]: + metastores: Dict[str, Metastore] = {} + assigned_metastore = self.unity_catalog_api_proxy.assigned_metastore() + if assigned_metastore: + metastores[assigned_metastore.metastore_id] = assigned_metastore for metastore in self.unity_catalog_api_proxy.metastores(): + metastores[metastore.metastore_id] = metastore + + for metastore in metastores.values(): if not self.config.metastore_id_pattern.allowed(metastore.metastore_id): self.report.metastores.dropped(metastore.metastore_id) continue diff --git a/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py b/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py index ecd05565bd..900f76f133 100644 --- a/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py +++ b/metadata-ingestion/tests/integration/unity/test_unity_catalog_ingest.py @@ -49,7 +49,24 @@ def register_mock_data(unity_catalog_api_instance): } ] } - + unity_catalog_api_instance.get_metastore_summary.return_value = { + "name": "acryl metastore", + "storage_root": "s3://db-9063-5b7e3b6d3736-s3-root-bucket/metastore/2c983545-d403-4f87-9063-5b7e3b6d3736", + "default_data_access_config_id": "9a9bacc4-cd82-409f-b55c-8ad1b2bde8da", + "storage_root_credential_id": "9a9bacc4-cd82-409f-b55c-8ad1b2bde8da", + "storage_root_credential_name": "2c983545-d403-4f87-9063-5b7e3b6d3736-data-access-config-1666185153576", + "delta_sharing_scope": "INTERNAL", + "owner": "abc@acryl.io", + "privilege_model_version": "1.0", + "region": "us-west-1", + "metastore_id": "2c983545-d403-4f87-9063-5b7e3b6d3736", + "created_at": 1666185153375, + "created_by": "abc@acryl.io", + "updated_at": 1666185154797, + "updated_by": "abc@acryl.io", + "cloud": "aws", + "global_metastore_id": "aws:us-west-1:2c983545-d403-4f87-9063-5b7e3b6d3736", + } unity_catalog_api_instance.list_catalogs.return_value = { "catalogs": [ { @@ -132,7 +149,6 @@ def register_mock_data(unity_catalog_api_instance): }, ] } - unity_catalog_api_instance.list_tables.return_value = { "tables": [ {