acryl-datahub 1.1.0rc4__py3-none-any.whl → 1.1.0.1__py3-none-any.whl

datahub/ingestion/source/unity/source.py

@@ -1,9 +1,12 @@
 import logging
 import re
+import time
 from concurrent.futures import ThreadPoolExecutor
 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 (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
@@ -78,6 +81,7 @@ from datahub.ingestion.source.unity.proxy_types import (
     Catalog,
     Column,
     CustomCatalogType,
+    HiveTableType,
     Metastore,
     Notebook,
     NotebookId,
@@ -87,8 +91,17 @@ from datahub.ingestion.source.unity.proxy_types import (
     TableReference,
 )
 from datahub.ingestion.source.unity.report import UnityCatalogReport
+from datahub.ingestion.source.unity.tag_entities import (
+    UnityCatalogTagPlatformResource,
+    UnityCatalogTagPlatformResourceId,
+)
 from datahub.ingestion.source.unity.usage import UnityCatalogUsageExtractor
-from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
+from datahub.metadata.com.linkedin.pegasus2avro.common import (
+    GlobalTags,
+    MetadataAttribution,
+    Siblings,
+    TagAssociation,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     DatasetLineageType,
     FineGrainedLineage,
@@ -116,6 +129,7 @@ from datahub.metadata.schema_classes import (
     UpstreamClass,
     UpstreamLineageClass,
 )
+from datahub.metadata.urns import TagUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import (
     SqlParsingResult,
@@ -162,6 +176,7 @@ 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
 
     def get_report(self) -> UnityCatalogReport:
         return self.report
@@ -211,6 +226,10 @@ 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
+            )
 
     def init_hive_metastore_proxy(self):
         self.hive_metastore_proxy: Optional[HiveMetastoreProxy] = None
@@ -506,13 +525,42 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         yield from self.add_table_to_dataset_container(dataset_urn, schema)
 
         table_props = self._create_table_property_aspect(table)
+        tags = None
+        if not isinstance(table.table_type, HiveTableType) and self.config.include_tags:
+            try:
+                table_tags = self._get_table_tags(
+                    table.ref.catalog, table.ref.schema, table.ref.table
+                )
+                if table_tags:
+                    logger.debug(f"Table tags for {table.ref}: {table_tags}")
+                    attribution = MetadataAttribution(
+                        # source="unity-catalog",
+                        actor="urn:li:corpuser:datahub",
+                        time=int(time.time() * 1000),
+                    )
+                    tags = GlobalTags(
+                        tags=[
+                            TagAssociation(
+                                tag=tag.to_datahub_tag_urn().urn(),
+                                attribution=attribution,
+                            )
+                            for tag in table_tags
+                        ]
+                    )
+
+                    yield from self.gen_platform_resources(table_tags)
+
+            except Exception as e:
+                logger.exception(f"Error fetching table {table.ref} tags", exc_info=e)
 
         view_props = None
         if table.view_definition:
             view_props = self._create_view_property_aspect(table)
 
         sub_type = self._create_table_sub_type_aspect(table)
-        schema_metadata = self._create_schema_metadata_aspect(table)
+        schema_metadata, platform_resources = self._create_schema_metadata_aspect(table)
+        yield from platform_resources
+
         domain = self._get_domain_aspect(dataset_name=table.ref.qualified_table_name)
         ownership = self._create_table_ownership_aspect(table)
         data_platform_instance = self._create_data_platform_instance_aspect()
@@ -585,6 +633,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                     domain,
                     data_platform_instance,
                     lineage,
+                    tags,
                 ],
             )
         ]
@@ -718,6 +767,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
     def gen_schema_containers(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
         domain_urn = self._gen_domain_urn(f"{schema.catalog.name}.{schema.name}")
+        schema_tags = self.unity_catalog_api_proxy.get_schema_tags(
+            schema.catalog.name
+        ).get(f"{schema.catalog.name}.{schema.name}", [])
+        if schema_tags:
+            logger.debug(f"Schema tags for {schema.name}: {schema_tags}")
+            # Generate platform resources for schema tags
+            yield from self.gen_platform_resources(schema_tags)
 
         schema_container_key = self.gen_schema_key(schema)
         yield from gen_containers(
@@ -729,6 +785,9 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             description=schema.comment,
             owner_urn=self.get_owner_urn(schema.owner),
             external_url=f"{self.external_url_base}/{schema.catalog.name}/{schema.name}",
+            tags=[tag.to_datahub_tag_urn().urn() for tag in schema_tags]
+            if schema_tags
+            else None,
         )
 
     def gen_metastore_containers(
@@ -749,6 +808,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
     def gen_catalog_containers(self, catalog: Catalog) -> Iterable[MetadataWorkUnit]:
         domain_urn = self._gen_domain_urn(catalog.name)
+        catalog_tags = self.unity_catalog_api_proxy.get_catalog_tags(catalog.name).get(
+            catalog.name, []
+        )
+        if catalog_tags:
+            logger.debug(f"Schema tags for {catalog.name}: {catalog_tags}")
+            # Generate platform resources for schema tags
+            yield from self.gen_platform_resources(catalog_tags)
 
         catalog_container_key = self.gen_catalog_key(catalog)
         yield from gen_containers(
@@ -764,6 +830,9 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             description=catalog.comment,
             owner_urn=self.get_owner_urn(catalog.owner),
             external_url=f"{self.external_url_base}/{catalog.name}",
+            tags=[tag.to_datahub_tag_urn().urn() for tag in catalog_tags]
+            if catalog_tags
+            else None,
         )
 
     def gen_schema_key(self, schema: Schema) -> ContainerKey:
@@ -832,6 +901,30 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             dataset_urn=dataset_urn,
         )
 
+    def _get_catalog_tags(
+        self, catalog: str, schema: str, table: str
+    ) -> List[UnityCatalogTag]:
+        all_tags = self.unity_catalog_api_proxy.get_catalog_tags(catalog)
+        return all_tags.get(f"{catalog}", [])
+
+    def _get_schema_tags(
+        self, catalog: str, schema: str, table: str
+    ) -> List[UnityCatalogTag]:
+        all_tags = self.unity_catalog_api_proxy.get_schema_tags(catalog)
+        return all_tags.get(f"{catalog}.{schema}", [])
+
+    def _get_table_tags(
+        self, catalog: str, schema: str, table: str
+    ) -> List[UnityCatalogTag]:
+        all_tags = self.unity_catalog_api_proxy.get_table_tags(catalog)
+        return all_tags.get(f"{catalog}.{schema}.{table}", [])
+
+    def _get_column_tags(
+        self, catalog: str, schema: str, table: str, column: str
+    ) -> List[UnityCatalogTag]:
+        all_tags = self.unity_catalog_api_proxy.get_column_tags(catalog)
+        return all_tags.get(f"{catalog}.{schema}.{table}.{column}", [])
+
     def _create_table_property_aspect(self, table: Table) -> DatasetPropertiesClass:
         custom_properties: dict = {}
         if table.storage_location is not None:
@@ -921,30 +1014,88 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             materialized=False, viewLanguage="SQL", viewLogic=table.view_definition
         )
 
-    def _create_schema_metadata_aspect(self, table: Table) -> SchemaMetadataClass:
-        schema_fields: List[SchemaFieldClass] = []
+    def gen_platform_resources(
+        self, tags: List[UnityCatalogTag]
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.ctx.graph and self.platform_resource_repository:
+            for tag in tags:
+                platform_resource_id = UnityCatalogTagPlatformResourceId.from_tag(
+                    platform_instance=self.platform_instance_name,
+                    platform_resource_repository=self.platform_resource_repository,
+                    tag=tag,
+                )
+                logger.debug(f"Created platform resource {platform_resource_id}")
 
+                unity_catalog_tag = UnityCatalogTagPlatformResource.get_from_datahub(
+                    platform_resource_id, self.platform_resource_repository, False
+                )
+                if (
+                    tag.to_datahub_tag_urn().urn()
+                    not in unity_catalog_tag.datahub_linked_resources().urns
+                ):
+                    unity_catalog_tag.datahub_linked_resources().add(
+                        tag.to_datahub_tag_urn().urn()
+                    )
+                    platform_resource = unity_catalog_tag.as_platform_resource()
+                    for mcp in platform_resource.to_mcps():
+                        yield MetadataWorkUnit(
+                            id=f"platform_resource-{platform_resource.id}",
+                            mcp=mcp,
+                        )
+
+    def _create_schema_metadata_aspect(
+        self, table: Table
+    ) -> Tuple[SchemaMetadataClass, Iterable[MetadataWorkUnit]]:
+        schema_fields: List[SchemaFieldClass] = []
+        unique_tags: Set[UnityCatalogTag] = set()
         for column in table.columns:
-            schema_fields.extend(self._create_schema_field(column))
-
-        return SchemaMetadataClass(
-            schemaName=table.id,
-            platform=make_data_platform_urn(self.platform),
-            fields=schema_fields,
-            hash="",
-            version=0,
-            platformSchema=MySqlDDLClass(tableSchema=""),
+            tag_urns: Optional[List[TagUrn]] = None
+            if self.config.include_tags:
+                column_tags = self._get_column_tags(
+                    table.ref.catalog, table.ref.schema, table.ref.table, column.name
+                )
+                unique_tags.update(column_tags)
+                tag_urns = [tag.to_datahub_tag_urn() for tag in column_tags]
+            schema_fields.extend(self._create_schema_field(column, tag_urns))
+
+        platform_resources = self.gen_platform_resources(list(unique_tags))
+        return (
+            SchemaMetadataClass(
+                schemaName=table.id,
+                platform=make_data_platform_urn(self.platform),
+                fields=schema_fields,
+                hash="",
+                version=0,
+                platformSchema=MySqlDDLClass(tableSchema=""),
+            ),
+            platform_resources,
         )
 
     @staticmethod
-    def _create_schema_field(column: Column) -> List[SchemaFieldClass]:
+    def _create_schema_field(
+        column: Column, tags: Optional[List[TagUrn]]
+    ) -> List[SchemaFieldClass]:
         _COMPLEX_TYPE = re.compile("^(struct|array)")
-
+        global_tags: Optional[GlobalTags] = None
         if _COMPLEX_TYPE.match(column.type_text.lower()):
             return get_schema_fields_for_hive_column(
                 column.name, column.type_text.lower(), description=column.comment
             )
         else:
+            if tags is not None:
+                logger.debug(f"Column tags are: {tags}")
+                attribution = MetadataAttribution(
+                    source="urn:li:dataPlatform:unity-catalog",
+                    actor="urn:li:corpuser:datahub",
+                    time=int(time.time() * 1000),
+                )
+                global_tags = GlobalTags(
+                    tags=[
+                        TagAssociation(tag=tag.urn(), attribution=attribution)
+                        for tag in tags
+                    ]
+                )
+
             return [
                 SchemaFieldClass(
                     fieldPath=column.name,
@@ -954,6 +1105,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                     nativeDataType=column.type_text,
                     nullable=column.nullable,
                     description=column.comment,
+                    globalTags=global_tags if tags else None,
                 )
             ]
 

datahub/ingestion/source/unity/tag_entities.py

@@ -0,0 +1,295 @@
+import logging
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+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
+
+
+class UnityCatalogTagSyncContext(BaseModel):
+    # it is intentionally empty
+    platform_instance: Optional[str] = None
+
+
+logger = logging.getLogger(__name__)
+
+
+class UnityCatalogTagPlatformResourceId(BaseModel, ExternalEntityId):
+    """
+    A SnowflakeTagId is a unique identifier for a Snowflake tag.
+    """
+
+    tag_key: str
+    tag_value: Optional[str] = None
+    platform_instance: Optional[str]
+    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:
+        return "UnityCatalogTagPlatformResource"
+
+    def to_platform_resource_key(self) -> PlatformResourceKey:
+        return PlatformResourceKey(
+            platform="databricks",
+            resource_type=str(UnityCatalogTagPlatformResourceId._RESOURCE_TYPE()),
+            primary_key=f"{self.tag_key}:{self.tag_value}",
+            platform_instance=self.platform_instance,
+        )
+
+    @classmethod
+    def from_tag(
+        cls,
+        tag: UnityCatalogTag,
+        platform_instance: Optional[str],
+        platform_resource_repository: PlatformResourceRepository,
+        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
+            ),
+        )
+        if existing_platform_resource:
+            logger.info(
+                f"Found existing UnityCatalogTagPlatformResourceId for tag {tag.key.original}: {existing_platform_resource}"
+            )
+            return existing_platform_resource
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=tag.key.original,
+            tag_value=tag.value.original if tag.value is not None else None,
+            platform_instance=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,
+        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
+        )
+        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
+                )
+            return new_unity_catalog_tag_id
+        raise ValueError(f"Unable to create SnowflakeTagId from DataHub URN: {urn}")
+
+    @classmethod
+    def generate_tag_id(
+        cls, graph: DataHubGraph, tag_sync_context: UnityCatalogTagSyncContext, urn: str
+    ) -> "UnityCatalogTagPlatformResourceId":
+        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
+                )
+            )
+        else:
+            raise ValueError(f"Unsupported entity type {entity_type} for URN {urn}")
+        return new_unity_catalog_tag_id
+
+    @classmethod
+    def from_datahub_tag(
+        cls, tag_urn: TagUrn, tag_sync_context: UnityCatalogTagSyncContext
+    ) -> "UnityCatalogTagPlatformResourceId":
+        uc_tag = UnityCatalogTag.from_urn(tag_urn)
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=str(uc_tag.key),
+            tag_value=str(uc_tag.value) if uc_tag.value is not None else None,
+            platform_instance=tag_sync_context.platform_instance,
+            exists_in_unity_catalog=False,
+        )
+
+
+class UnityCatalogTagPlatformResource(BaseModel, ExternalEntity):
+    datahub_urns: LinkedResourceSet
+    managed_by_datahub: bool
+    id: UnityCatalogTagPlatformResourceId
+    allowed_values: Optional[List[str]]
+
+    def get_id(self) -> ExternalEntityId:
+        return self.id
+
+    def is_managed_by_datahub(self) -> bool:
+        return self.managed_by_datahub
+
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        return self.datahub_urns
+
+    def as_platform_resource(self) -> PlatformResource:
+        return PlatformResource.create(
+            key=self.id.to_platform_resource_key(),
+            secondary_keys=[u for u in self.datahub_urns.urns],
+            value=self,
+        )
+
+    @classmethod
+    def get_from_datahub(
+        cls,
+        unity_catalog_tag_id: UnityCatalogTagPlatformResourceId,
+        platform_resource_repository: PlatformResourceRepository,
+        managed_by_datahub: bool = False,
+    ) -> "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
+        return cls(
+            id=unity_catalog_tag_id,
+            datahub_urns=LinkedResourceSet(urns=[]),
+            managed_by_datahub=managed_by_datahub,
+            allowed_values=None,
+        )