acryl-datahub 1.2.0.4rc3__py3-none-any.whl → 1.2.0.5__py3-none-any.whl

datahub/ingestion/source/unity/proxy.py

@@ -141,6 +141,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         report: UnityCatalogReport,
         hive_metastore_proxy: Optional[HiveMetastoreProxy] = None,
         lineage_data_source: LineageDataSource = LineageDataSource.AUTO,
+        databricks_api_page_size: int = 0,
     ):
         self._workspace_client = WorkspaceClient(
             host=workspace_url,
@@ -152,6 +153,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
         self.lineage_data_source = lineage_data_source
+        self.databricks_api_page_size = databricks_api_page_size
         self._sql_connection_params = {
             "server_hostname": self._workspace_client.config.host.replace(
                 "https://", ""
@@ -161,7 +163,11 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         }
 
     def check_basic_connectivity(self) -> bool:
-        return bool(self._workspace_client.catalogs.list(include_browse=True))
+        return bool(
+            self._workspace_client.catalogs.list(
+                include_browse=True, max_results=self.databricks_api_page_size
+            )
+        )
 
     def assigned_metastore(self) -> Optional[Metastore]:
         response = self._workspace_client.metastores.summary()
@@ -171,7 +177,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         if self.hive_metastore_proxy:
             yield self.hive_metastore_proxy.hive_metastore_catalog(metastore)
 
-        response = self._workspace_client.catalogs.list(include_browse=True)
+        response = self._workspace_client.catalogs.list(
+            include_browse=True, max_results=self.databricks_api_page_size
+        )
         if not response:
             logger.info("Catalogs not found")
             return
@@ -203,7 +211,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             yield from self.hive_metastore_proxy.hive_metastore_schemas(catalog)
             return
         response = self._workspace_client.schemas.list(
-            catalog_name=catalog.name, include_browse=True
+            catalog_name=catalog.name,
+            include_browse=True,
+            max_results=self.databricks_api_page_size,
         )
         if not response:
             logger.info(f"Schemas not found for catalog {catalog.id}")
@@ -225,6 +235,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 catalog_name=schema.catalog.name,
                 schema_name=schema.name,
                 include_browse=True,
+                max_results=self.databricks_api_page_size,
             )
             if not response:
                 logger.info(f"Tables not found for schema {schema.id}")
@@ -257,7 +268,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         return group_list
 
     def workspace_notebooks(self) -> Iterable[Notebook]:
-        for obj in self._workspace_client.workspace.list("/", recursive=True):
+        workspace_objects_iter = self._workspace_client.workspace.list(
+            "/", recursive=True, max_results=self.databricks_api_page_size
+        )
+        for obj in workspace_objects_iter:
             if obj.object_type == ObjectType.NOTEBOOK and obj.object_id and obj.path:
                 yield Notebook(
                     id=obj.object_id,
@@ -299,7 +313,6 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     def _query_history(
         self,
         filter_by: QueryFilterWithStatementTypes,
-        max_results: int = 1000,
         include_metrics: bool = False,
     ) -> Iterable[QueryInfo]:
         """Manual implementation of the query_history.list() endpoint.
@@ -311,9 +324,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         """
         method = "GET"
         path = "/api/2.0/sql/history/queries"
+
         body: Dict[str, Any] = {
             "include_metrics": include_metrics,
-            "max_results": max_results,  # Max batch size
+            "max_results": self.databricks_api_page_size,  # Max batch size
         }
 
         response: dict = self._workspace_client.api_client.do(  # type: ignore

datahub/ingestion/source/unity/report.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from typing import TYPE_CHECKING, Optional, Tuple
 
 from datahub.ingestion.api.report import EntityFilterReport, Report
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
@@ -7,6 +7,11 @@ from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.perf_timer import PerfTimer
 
+if TYPE_CHECKING:
+    from datahub.ingestion.source.unity.platform_resource_repository import (
+        UnityCatalogPlatformResourceRepository,
+    )
+
 
 @dataclass
 class UnityCatalogUsagePerfReport(Report):
@@ -61,3 +66,6 @@ class UnityCatalogReport(IngestionStageReport, SQLSourceReport):
     num_tables_missing_name: int = 0
     num_columns_missing_name: int = 0
     num_queries_missing_info: int = 0
+
+    # Platform resource repository for automatic cache statistics via SupportsAsObj
+    tag_urn_resolver_cache: Optional["UnityCatalogPlatformResourceRepository"] = None

datahub/ingestion/source/unity/source.py

@@ -4,7 +4,6 @@ import time
 from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 from urllib.parse import urljoin
 
-from datahub.api.entities.external.external_entities import PlatformResourceRepository
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import (
     UNKNOWN_USER,
@@ -77,6 +76,9 @@ from datahub.ingestion.source.unity.hive_metastore_proxy import (
     HIVE_METASTORE,
     HiveMetastoreProxy,
 )
+from datahub.ingestion.source.unity.platform_resource_repository import (
+    UnityCatalogPlatformResourceRepository,
+)
 from datahub.ingestion.source.unity.proxy import UnityCatalogApiProxy
 from datahub.ingestion.source.unity.proxy_types import (
     DATA_TYPE_REGISTRY,
@@ -187,7 +189,9 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     platform: str = "databricks"
     platform_instance_name: Optional[str]
     sql_parser_schema_resolver: Optional[SchemaResolver] = None
-    platform_resource_repository: Optional[PlatformResourceRepository] = None
+    platform_resource_repository: Optional[UnityCatalogPlatformResourceRepository] = (
+        None
+    )
 
     def get_report(self) -> UnityCatalogReport:
         return self.report
@@ -207,6 +211,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
             lineage_data_source=config.lineage_data_source,
+            databricks_api_page_size=config.databricks_api_page_size,
         )
 
         self.external_url_base = urljoin(self.config.workspace_url, "/explore/data")
@@ -239,9 +244,15 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         # Global map of tables, for profiling
         self.tables: FileBackedDict[Table] = FileBackedDict()
         if self.ctx.graph:
-            self.platform_resource_repository = PlatformResourceRepository(
-                self.ctx.graph
+            self.platform_resource_repository = UnityCatalogPlatformResourceRepository(
+                self.ctx.graph, platform_instance=self.platform_instance_name
             )
+        else:
+            self.platform_resource_repository = None
+
+        # Include platform resource repository in report for automatic cache statistics
+        if self.config.include_tags and self.platform_resource_repository:
+            self.report.tag_urn_resolver_cache = self.platform_resource_repository
 
     def init_hive_metastore_proxy(self):
         self.hive_metastore_proxy: Optional[HiveMetastoreProxy] = None
@@ -1079,25 +1090,49 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             materialized=False, viewLanguage="SQL", viewLogic=table.view_definition
         )
 
+    def get_or_create_from_unity_tag(
+        self,
+        unity_tag: UnityCatalogTag,
+        platform_instance: Optional[str],
+        managed_by_datahub: bool = False,
+    ) -> UnityCatalogTagPlatformResource:
+        """
+        Optimized helper to get or create a Unity Catalog tag platform resource.
+        This eliminates the duplicate search by skipping the from_tag method which was
+        doing a redundant search before get_from_datahub.
+        """
+        # Create the platform resource ID directly without the from_tag search
+        platform_resource_id = UnityCatalogTagPlatformResourceId(
+            tag_key=unity_tag.key.raw_text,
+            tag_value=unity_tag.value.raw_text if unity_tag.value is not None else None,
+            platform_instance=platform_instance,
+            exists_in_unity_catalog=True,  # We got it from Unity Catalog
+            persisted=False,
+        )
+
+        # Use the repository's get_entity_from_datahub method which handles
+        # searching and caching internally - this is the ONLY search we need
+        if self.platform_resource_repository is None:
+            raise ValueError("Platform resource repository not initialized")
+        return self.platform_resource_repository.get_entity_from_datahub(
+            platform_resource_id,
+            managed_by_datahub,
+        )
+
     def gen_platform_resources(
         self, tags: List[UnityCatalogTag]
     ) -> Iterable[MetadataWorkUnit]:
         if self.ctx.graph and self.platform_resource_repository:
             for tag in tags:
                 try:
-                    platform_resource_id = UnityCatalogTagPlatformResourceId.from_tag(
-                        platform_instance=self.platform_instance_name,
-                        platform_resource_repository=self.platform_resource_repository,
-                        tag=tag,
+                    # Use optimized helper method that combines ID creation and entity retrieval
+                    unity_catalog_tag = self.get_or_create_from_unity_tag(
+                        tag,
+                        self.platform_instance_name,
+                        managed_by_datahub=False,
                     )
-                    logger.debug(f"Created platform resource {platform_resource_id}")
-
-                    unity_catalog_tag = (
-                        UnityCatalogTagPlatformResource.get_from_datahub(
-                            platform_resource_id,
-                            self.platform_resource_repository,
-                            False,
-                        )
+                    logger.debug(
+                        f"Retrieved/created platform resource for tag {tag.key.raw_text}"
                     )
                     if (
                         tag.to_datahub_tag_urn().urn()

datahub/ingestion/source/unity/tag_entities.py

@@ -1,5 +1,10 @@
 import logging
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.source.unity.platform_resource_repository import (
+        UnityCatalogPlatformResourceRepository,
+    )
 
 from pydantic import BaseModel
 
@@ -7,17 +12,14 @@ from datahub.api.entities.external.external_entities import (
     ExternalEntity,
     ExternalEntityId,
     LinkedResourceSet,
-    PlatformResourceRepository,
 )
 from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.api.entities.platformresource.platform_resource import (
     PlatformResource,
     PlatformResourceKey,
-    PlatformResourceSearchFields,
 )
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import TagUrn
-from datahub.utilities.search_utils import ElasticDocumentQuery
 from datahub.utilities.urns.urn import Urn
 
 
@@ -29,9 +31,9 @@ class UnityCatalogTagSyncContext(BaseModel):
 logger = logging.getLogger(__name__)
 
 
-class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
+class UnityCatalogTagPlatformResourceId(ExternalEntityId):
     """
-    A SnowflakeTagId is a unique identifier for a Snowflake tag.
+    A Unity Catalog tag platform resource ID.
     """
 
     tag_key: str
@@ -40,9 +42,6 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
     exists_in_unity_catalog: bool = False
     persisted: bool = False
 
-    def __hash__(self) -> int:
-        return hash(self.to_platform_resource_key().id)
-
     # this is a hack to make sure the property is a string and not private pydantic field
     @staticmethod
     def _RESOURCE_TYPE() -> str:
@@ -57,26 +56,21 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
         )
 
     @classmethod
-    def from_tag(
+    def get_or_create_from_tag(
         cls,
         tag: UnityCatalogTag,
-        platform_instance: Optional[str],
-        platform_resource_repository: PlatformResourceRepository,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
         exists_in_unity_catalog: bool = False,
     ) -> "UnityCatalogTagPlatformResourceId":
         """
         Creates a UnityCatalogTagPlatformResourceId from a UnityCatalogTag.
         """
 
-        existing_platform_resource = cls.search_by_urn(
-            tag.to_datahub_tag_urn().urn(),
-            platform_resource_repository=platform_resource_repository,
-            tag_sync_context=UnityCatalogTagSyncContext(
-                platform_instance=platform_instance
-            ),
+        existing_platform_resource = platform_resource_repository.search_entity_by_urn(
+            tag.to_datahub_tag_urn().urn()
         )
         if existing_platform_resource:
-            logger.info(
+            logger.debug(
                 f"Found existing UnityCatalogTagPlatformResourceId for tag {tag.key.raw_text}: {existing_platform_resource}"
             )
             return existing_platform_resource
@@ -84,105 +78,46 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
         return UnityCatalogTagPlatformResourceId(
             tag_key=tag.key.raw_text,
             tag_value=tag.value.raw_text if tag.value is not None else None,
-            platform_instance=platform_instance,
+            platform_instance=platform_resource_repository.platform_instance,
             exists_in_unity_catalog=exists_in_unity_catalog,
             persisted=False,
         )
 
-    @classmethod
-    def search_by_urn(
-        cls,
-        urn: str,
-        platform_resource_repository: PlatformResourceRepository,
-        tag_sync_context: UnityCatalogTagSyncContext,
-    ) -> Optional["UnityCatalogTagPlatformResourceId"]:
-        mapped_tags = [
-            t
-            for t in platform_resource_repository.search_by_filter(
-                ElasticDocumentQuery.create_from(
-                    (
-                        PlatformResourceSearchFields.RESOURCE_TYPE,
-                        str(UnityCatalogTagPlatformResourceId._RESOURCE_TYPE()),
-                    ),
-                    (PlatformResourceSearchFields.SECONDARY_KEYS, urn),
-                )
-            )
-        ]
-        logger.info(
-            f"Found {len(mapped_tags)} mapped tags for URN {urn}. {mapped_tags}"
-        )
-        if len(mapped_tags) > 0:
-            for platform_resource in mapped_tags:
-                if (
-                    platform_resource.resource_info
-                    and platform_resource.resource_info.value
-                ):
-                    unity_catalog_tag = UnityCatalogTagPlatformResource(
-                        **platform_resource.resource_info.value.as_pydantic_object(
-                            UnityCatalogTagPlatformResource
-                        ).dict()
-                    )
-                    if (
-                        unity_catalog_tag.id.platform_instance
-                        == tag_sync_context.platform_instance
-                    ):
-                        unity_catalog_tag_id = unity_catalog_tag.id
-                        unity_catalog_tag_id.exists_in_unity_catalog = True
-                        unity_catalog_tag_id.persisted = True
-                        return unity_catalog_tag_id
-                else:
-                    logger.warning(
-                        f"Platform resource {platform_resource} does not have a resource_info value"
-                    )
-                    continue
-
-            # If we reach here, it means we did not find a mapped tag for the URN
-            logger.info(
-                f"No mapped tag found for URN {urn} with platform instance {tag_sync_context.platform_instance}. Creating a new UnityCatalogTagPlatformResourceId."
-            )
-        return None
-
     @classmethod
     def from_datahub_urn(
         cls,
         urn: str,
-        platform_resource_repository: PlatformResourceRepository,
         tag_sync_context: UnityCatalogTagSyncContext,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
         graph: DataHubGraph,
     ) -> "UnityCatalogTagPlatformResourceId":
         """
         Creates a UnityCatalogTagPlatformResourceId from a DataHub URN.
         """
-        # First we check if we already have a mapped platform resource for this
-        # urn that is of the type UnityCatalogTagPlatformResource
-        # If we do, we can use it to create the UnityCatalogTagPlatformResourceId
-        # Else, we need to generate a new UnityCatalogTagPlatformResourceId
-        existing_platform_resource_id = cls.search_by_urn(
-            urn, platform_resource_repository, tag_sync_context
+        existing_platform_resource_id = (
+            platform_resource_repository.search_entity_by_urn(urn)
         )
         if existing_platform_resource_id:
-            logger.info(
-                f"Found existing UnityCatalogTagPlatformResourceId for URN {urn}: {existing_platform_resource_id}"
-            )
             return existing_platform_resource_id
 
-        # Otherwise, we need to create a new UnityCatalogTagPlatformResourceId
         new_unity_catalog_tag_id = cls.generate_tag_id(graph, tag_sync_context, urn)
         if new_unity_catalog_tag_id:
-            # we then check if this tag has already been ingested as a platform
-            # resource in the platform resource repository
             resource_key = platform_resource_repository.get(
                 new_unity_catalog_tag_id.to_platform_resource_key()
             )
             if resource_key:
-                logger.info(
-                    f"Tag {new_unity_catalog_tag_id} already exists in platform resource repository with {resource_key}"
-                )
-                new_unity_catalog_tag_id.exists_in_unity_catalog = (
-                    True  # TODO: Check if this is a safe assumption
+                # Create a new ID with the correct state instead of mutating
+                return UnityCatalogTagPlatformResourceId(
+                    tag_key=new_unity_catalog_tag_id.tag_key,
+                    tag_value=new_unity_catalog_tag_id.tag_value,
+                    platform_instance=new_unity_catalog_tag_id.platform_instance,
+                    exists_in_unity_catalog=True,  # This tag exists in Unity Catalog
+                    persisted=new_unity_catalog_tag_id.persisted,
                 )
             return new_unity_catalog_tag_id
-        raise ValueError(f"Unable to create SnowflakeTagId from DataHub URN: {urn}")
+        raise ValueError(
+            f"Unable to create Unity Catalog tag ID from DataHub URN: {urn}"
+        )
 
     @classmethod
     def generate_tag_id(
@@ -191,14 +126,11 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
         parsed_urn = Urn.from_string(urn)
         entity_type = parsed_urn.entity_type
         if entity_type == "tag":
-            new_unity_catalog_tag_id = (
-                UnityCatalogTagPlatformResourceId.from_datahub_tag(
-                    TagUrn.from_string(urn), tag_sync_context
-                )
+            return UnityCatalogTagPlatformResourceId.from_datahub_tag(
+                TagUrn.from_string(urn), tag_sync_context
             )
         else:
             raise ValueError(f"Unsupported entity type {entity_type} for URN {urn}")
-        return new_unity_catalog_tag_id
 
     @classmethod
     def from_datahub_tag(
@@ -214,7 +146,7 @@ class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
         )
 
 
-class UnityCatalogTagPlatformResource(BaseModel, ExternalEntity):
+class UnityCatalogTagPlatformResource(ExternalEntity):
     datahub_urns: LinkedResourceSet
     managed_by_datahub: bool
     id: UnityCatalogTagPlatformResourceId
@@ -237,58 +169,28 @@ class UnityCatalogTagPlatformResource(BaseModel, ExternalEntity):
         )
 
     @classmethod
-    def get_from_datahub(
+    def create_default(
         cls,
-        unity_catalog_tag_id: UnityCatalogTagPlatformResourceId,
-        platform_resource_repository: PlatformResourceRepository,
-        managed_by_datahub: bool = False,
+        entity_id: ExternalEntityId,
+        managed_by_datahub: bool,
     ) -> "UnityCatalogTagPlatformResource":
-        # Search for linked DataHub URNs
-        platform_resources = [
-            r
-            for r in platform_resource_repository.search_by_filter(
-                ElasticDocumentQuery.create_from(
-                    (
-                        PlatformResourceSearchFields.RESOURCE_TYPE,
-                        str(UnityCatalogTagPlatformResourceId._RESOURCE_TYPE()),
-                    ),
-                    (
-                        PlatformResourceSearchFields.PRIMARY_KEY,
-                        f"{unity_catalog_tag_id.tag_key}/{unity_catalog_tag_id.tag_value}",
-                    ),
-                )
-            )
-        ]
-        if len(platform_resources) == 1:
-            platform_resource: PlatformResource = platform_resources[0]
-            if (
-                platform_resource.resource_info
-                and platform_resource.resource_info.value
-            ):
-                unity_catalog_tag = UnityCatalogTagPlatformResource(
-                    **platform_resource.resource_info.value.as_pydantic_object(
-                        UnityCatalogTagPlatformResource
-                    ).dict()
-                )
-                return unity_catalog_tag
-        else:
-            for platform_resource in platform_resources:
-                if (
-                    platform_resource.resource_info
-                    and platform_resource.resource_info.value
-                ):
-                    unity_catalog_tag = UnityCatalogTagPlatformResource(
-                        **platform_resource.resource_info.value.as_pydantic_object(
-                            UnityCatalogTagPlatformResource
-                        ).dict()
-                    )
-                    if (
-                        unity_catalog_tag.id.platform_instance
-                        == unity_catalog_tag_id.platform_instance
-                    ):
-                        return unity_catalog_tag
+        """Create a default Unity Catalog tag entity when none found in DataHub."""
+        # Type narrowing: we know this will be a UnityCatalogTagPlatformResourceId
+        assert isinstance(entity_id, UnityCatalogTagPlatformResourceId), (
+            f"Expected UnityCatalogTagPlatformResourceId, got {type(entity_id)}"
+        )
+
+        # Create a new entity ID with correct default state instead of mutating
+        default_entity_id = UnityCatalogTagPlatformResourceId(
+            tag_key=entity_id.tag_key,
+            tag_value=entity_id.tag_value,
+            platform_instance=entity_id.platform_instance,
+            exists_in_unity_catalog=False,  # New entities don't exist in Unity Catalog yet
+            persisted=False,  # New entities are not persisted yet
+        )
+
         return cls(
-            id=unity_catalog_tag_id,
+            id=default_entity_id,
             datahub_urns=LinkedResourceSet(urns=[]),
             managed_by_datahub=managed_by_datahub,
             allowed_values=None,