acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/unity/source.py

@@ -1,16 +1,20 @@
 import logging
 import re
-from concurrent.futures import ThreadPoolExecutor
+import time
 from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
 from urllib.parse import urljoin
 
+from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import (
+    UNKNOWN_USER,
     make_data_platform_urn,
     make_dataplatform_instance_urn,
     make_dataset_urn_with_platform_instance,
     make_domain_urn,
     make_group_urn,
+    make_ml_model_group_urn,
     make_schema_field_urn,
+    make_ts_millis,
     make_user_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
@@ -23,9 +27,9 @@ from datahub.emitter.mcp_builder import (
     UnitySchemaKey,
     UnitySchemaKeyWithMetastore,
     add_dataset_to_container,
+    add_entity_to_container,
     gen_containers,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -53,6 +57,7 @@ from datahub.ingestion.source.aws.s3_util import (
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
+    SourceCapabilityModifier,
 )
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -72,13 +77,19 @@ 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,
     Catalog,
     Column,
     CustomCatalogType,
+    HiveTableType,
     Metastore,
+    Model,
+    ModelVersion,
     Notebook,
     NotebookId,
     Schema,
@@ -87,8 +98,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,
@@ -98,11 +118,13 @@ from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     ViewProperties,
 )
 from datahub.metadata.schema_classes import (
+    AuditStampClass,
     BrowsePathsClass,
     DataPlatformInstanceClass,
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
     DomainsClass,
+    MLModelPropertiesClass,
     MySqlDDLClass,
     NullTypeClass,
     OwnerClass,
@@ -116,12 +138,10 @@ from datahub.metadata.schema_classes import (
     UpstreamClass,
     UpstreamLineageClass,
 )
+from datahub.metadata.urns import MlModelGroupUrn, MlModelUrn, TagUrn
+from datahub.sdk import MLModel, MLModelGroup
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import (
-    SqlParsingResult,
-    sqlglot_lineage,
-    view_definition_lineage_helper,
-)
+from datahub.sql_parsing.sql_parsing_aggregator import SqlParsingAggregator
 from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -138,23 +158,32 @@ logger: logging.Logger = logging.getLogger(__name__)
 @capability(SourceCapability.USAGE_STATS, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.CATALOG,
+        SourceCapabilityModifier.SCHEMA,
+    ],
+)
 @capability(SourceCapability.OWNERSHIP, "Supported via the `include_ownership` config")
 @capability(
     SourceCapability.DATA_PROFILING, "Supported via the `profiling.enabled` config"
 )
 @capability(
     SourceCapability.DELETION_DETECTION,
-    "Optionally enabled via `stateful_ingestion.remove_stale_metadata`",
+    "Enabled by default via stateful ingestion",
     supported=True,
 )
-@support_status(SupportStatus.INCUBATING)
+@capability(SourceCapability.TEST_CONNECTION, "Enabled by default")
+@support_status(SupportStatus.CERTIFIED)
 class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     """
     This plugin extracts the following metadata from Databricks Unity Catalog:
     - metastores
     - schemas
     - tables and column lineage
+    - model and model versions
     """
 
     config: UnityCatalogSourceConfig
@@ -162,6 +191,10 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     platform: str = "databricks"
     platform_instance_name: Optional[str]
     sql_parser_schema_resolver: Optional[SchemaResolver] = None
+    platform_resource_repository: Optional[UnityCatalogPlatformResourceRepository] = (
+        None
+    )
+    sql_parsing_aggregator: Optional[SqlParsingAggregator] = None
 
     def get_report(self) -> UnityCatalogReport:
         return self.report
@@ -180,6 +213,9 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             config.warehouse_id,
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
+            lineage_data_source=config.lineage_data_source,
+            usage_data_source=config.usage_data_source,
+            databricks_api_page_size=config.databricks_api_page_size,
         )
 
         self.external_url_base = urljoin(self.config.workspace_url, "/explore/data")
@@ -205,12 +241,31 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         self.table_refs: Set[TableReference] = set()
         self.view_refs: Set[TableReference] = set()
         self.notebooks: FileBackedDict[Notebook] = FileBackedDict()
-        self.view_definitions: FileBackedDict[Tuple[TableReference, str]] = (
-            FileBackedDict()
-        )
 
         # Global map of tables, for profiling
         self.tables: FileBackedDict[Table] = FileBackedDict()
+        if self.ctx.graph:
+            self.platform_resource_repository = UnityCatalogPlatformResourceRepository(
+                self.ctx.graph, platform_instance=self.platform_instance_name
+            )
+        else:
+            self.platform_resource_repository = None
+
+        if self.config._forced_disable_tag_extraction:
+            self.report.report_warning(
+                "Some features disabled because of configuration conflicts",
+                "Tag Extraction is disabled due to missing warehouse_id in config",
+            )
+
+        if self.config._forced_disable_hive_metastore_extraction:
+            self.report.report_warning(
+                "Some features disabled because of configuration conflicts",
+                "Hive Metastore Extraction is disabled due to missing warehouse_id in config",
+            )
+
+        # 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
@@ -229,6 +284,17 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                         platform_instance=self.config.platform_instance,
                         env=self.config.env,
                     )
+                    self.sql_parsing_aggregator = SqlParsingAggregator(
+                        platform=self.platform,
+                        platform_instance=self.config.platform_instance,
+                        env=self.config.env,
+                        schema_resolver=self.sql_parser_schema_resolver,
+                        generate_lineage=True,
+                        generate_queries=False,
+                        generate_usage_statistics=False,
+                        generate_operations=False,
+                    )
+                    self.report.sql_aggregator = self.sql_parsing_aggregator.report
             except Exception as e:
                 logger.debug("Exception", exc_info=True)
                 self.warn(
@@ -383,12 +449,12 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                         self.config.workspace_url, f"#notebook/{notebook.id}"
                     ),
                     created=(
-                        TimeStampClass(int(notebook.created_at.timestamp() * 1000))
+                        TimeStampClass(make_ts_millis(notebook.created_at))
                         if notebook.created_at
                         else None
                     ),
                     lastModified=(
-                        TimeStampClass(int(notebook.modified_at.timestamp() * 1000))
+                        TimeStampClass(make_ts_millis(notebook.modified_at))
                         if notebook.modified_at
                         else None
                     ),
@@ -407,17 +473,20 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         if not notebook.upstreams:
             return None
 
+        upstreams = []
+        for upstream_ref in notebook.upstreams:
+            timestamp = make_ts_millis(upstream_ref.last_updated)
+            upstreams.append(
+                self._create_upstream_class(
+                    self.gen_dataset_urn(upstream_ref),
+                    DatasetLineageTypeClass.COPY,
+                    timestamp,
+                )
+            )
+
         return MetadataChangeProposalWrapper(
             entityUrn=self.gen_notebook_urn(notebook),
-            aspect=UpstreamLineageClass(
-                upstreams=[
-                    UpstreamClass(
-                        dataset=self.gen_dataset_urn(upstream_ref),
-                        type=DatasetLineageTypeClass.COPY,
-                    )
-                    for upstream_ref in notebook.upstreams
-                ]
-            ),
+            aspect=UpstreamLineageClass(upstreams=upstreams),
         ).as_workunit()
 
     def process_metastores(self) -> Iterable[MetadataWorkUnit]:
@@ -436,14 +505,15 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         self, metastore: Optional[Metastore]
     ) -> Iterable[MetadataWorkUnit]:
         for catalog in self._get_catalogs(metastore):
-            if not self.config.catalog_pattern.allowed(catalog.id):
-                self.report.catalogs.dropped(catalog.id)
-                continue
+            with self.report.new_stage(f"Ingest catalog {catalog.id}"):
+                if not self.config.catalog_pattern.allowed(catalog.id):
+                    self.report.catalogs.dropped(catalog.id)
+                    continue
 
-            yield from self.gen_catalog_containers(catalog)
-            yield from self.process_schemas(catalog)
+                yield from self.gen_catalog_containers(catalog)
+                yield from self.process_schemas(catalog)
 
-            self.report.catalogs.processed(catalog.id)
+                self.report.catalogs.processed(catalog.id)
 
     def _get_catalogs(self, metastore: Optional[Metastore]) -> Iterable[Catalog]:
         if self.config.catalogs:
@@ -466,6 +536,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 yield from self.gen_schema_containers(schema)
                 try:
                     yield from self.process_tables(schema)
+                    yield from self.process_ml_models(schema)
                 except Exception as e:
                     logger.exception(f"Error parsing schema {schema}")
                     self.report.report_warning(
@@ -506,13 +577,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()
@@ -534,8 +634,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             self.sql_parser_schema_resolver.add_schema_metadata(
                 dataset_urn, schema_metadata
             )
-            if table.view_definition:
-                self.view_definitions[dataset_urn] = (table.ref, table.view_definition)
+            if table.view_definition and self.sql_parsing_aggregator:
+                self.sql_parsing_aggregator.add_view_definition(
+                    view_urn=dataset_urn,
+                    view_definition=table.view_definition,
+                    default_db=table.ref.catalog,
+                    default_schema=table.ref.schema,
+                )
 
         if (
             table_props.customProperties.get("table_type")
@@ -585,29 +690,107 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                     domain,
                     data_platform_instance,
                     lineage,
+                    tags,
                 ],
             )
         ]
 
+    def process_ml_models(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
+        for ml_model in self.unity_catalog_api_proxy.ml_models(
+            schema=schema, max_results=self.config.ml_model_max_results
+        ):
+            yield from self.process_ml_model(ml_model, schema)
+            ml_model_urn = self.gen_ml_model_urn(ml_model.id)
+            for ml_model_version in self.unity_catalog_api_proxy.ml_model_versions(
+                ml_model, include_aliases=self.config.include_ml_model_aliases
+            ):
+                yield from self.process_ml_model_version(
+                    ml_model_urn, ml_model_version, schema
+                )
+
+    def process_ml_model(
+        self, ml_model: Model, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        ml_model_group = MLModelGroup(
+            id=ml_model.id,
+            name=ml_model.name,
+            platform=self.platform,
+            platform_instance=schema.name,
+            env=self.config.env,
+            description=ml_model.description,
+            created=ml_model.created_at,
+            last_modified=ml_model.updated_at,
+        )
+        yield from ml_model_group.as_workunits()
+        yield from self.add_model_to_schema_container(str(ml_model_group.urn), schema)
+        self.report.ml_models.processed(ml_model.id)
+
+    def process_ml_model_version(
+        self, ml_model_urn: str, ml_model_version: ModelVersion, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        extra_aspects = []
+        if ml_model_version.created_at is not None:
+            created_time = int(ml_model_version.created_at.timestamp() * 1000)
+            created_actor = (
+                f"urn:li:platformResource:{ml_model_version.created_by}"
+                if ml_model_version.created_by
+                else None
+            )
+            extra_aspects.append(
+                MLModelPropertiesClass(
+                    created=TimeStampClass(time=created_time, actor=created_actor),
+                )
+            )
+
+        ml_model = MLModel(
+            id=ml_model_version.id,
+            name=ml_model_version.name,
+            version=str(ml_model_version.version),
+            aliases=ml_model_version.aliases,
+            description=ml_model_version.description,
+            model_group=ml_model_urn,
+            platform=self.platform,
+            last_modified=ml_model_version.updated_at,
+            extra_aspects=extra_aspects,
+        )
+
+        yield from ml_model.as_workunits()
+        yield from self.add_model_version_to_schema_container(str(ml_model.urn), schema)
+        self.report.ml_model_versions.processed(ml_model_version.id)
+
     def ingest_lineage(self, table: Table) -> Optional[UpstreamLineageClass]:
+        # Calculate datetime filters for lineage
+        lineage_start_time = None
+        lineage_end_time = self.config.end_time
+
+        if self.config.ignore_start_time_lineage:
+            lineage_start_time = None  # Ignore start time to get all lineage
+        else:
+            lineage_start_time = self.config.start_time
+
         if self.config.include_table_lineage:
             self.unity_catalog_api_proxy.table_lineage(
-                table, include_entity_lineage=self.config.include_notebooks
+                table,
+                include_entity_lineage=self.config.include_notebooks,
+                start_time=lineage_start_time,
+                end_time=lineage_end_time,
             )
 
         if self.config.include_column_lineage and table.upstreams:
             if len(table.columns) > self.config.column_lineage_column_limit:
                 self.report.num_column_lineage_skipped_column_count += 1
 
-            with ThreadPoolExecutor(
-                max_workers=self.config.lineage_max_workers
-            ) as executor:
-                for column in table.columns[: self.config.column_lineage_column_limit]:
-                    executor.submit(
-                        self.unity_catalog_api_proxy.get_column_lineage,
-                        table,
-                        column.name,
-                    )
+            column_names = [
+                column.name
+                for column in table.columns[: self.config.column_lineage_column_limit]
+            ]
+            self.unity_catalog_api_proxy.get_column_lineage(
+                table,
+                column_names,
+                max_workers=self.config.lineage_max_workers,
+                start_time=lineage_start_time,
+                end_time=lineage_end_time,
+            )
 
         return self._generate_lineage_aspect(self.gen_dataset_urn(table.ref), table)
 
@@ -635,18 +818,22 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 for d_col, u_cols in sorted(downstream_to_upstream_cols.items())
             )
 
+            timestamp = make_ts_millis(upstream_ref.last_updated)
             upstreams.append(
-                UpstreamClass(
-                    dataset=upstream_urn,
-                    type=DatasetLineageTypeClass.TRANSFORMED,
+                self._create_upstream_class(
+                    upstream_urn,
+                    DatasetLineageTypeClass.TRANSFORMED,
+                    timestamp,
                 )
             )
 
-        for notebook in table.upstream_notebooks:
+        for notebook in table.upstream_notebooks.values():
+            timestamp = make_ts_millis(notebook.last_updated)
             upstreams.append(
-                UpstreamClass(
-                    dataset=self.gen_notebook_urn(notebook),
-                    type=DatasetLineageTypeClass.TRANSFORMED,
+                self._create_upstream_class(
+                    self.gen_notebook_urn(notebook.id),
+                    DatasetLineageTypeClass.TRANSFORMED,
+                    timestamp,
                 )
             )
 
@@ -708,6 +895,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             env=self.config.env,
         )
 
+    def gen_ml_model_urn(self, name: str) -> str:
+        return make_ml_model_group_urn(
+            platform=self.platform,
+            group_name=name,
+            env=self.config.env,
+        )
+
     def gen_notebook_urn(self, notebook: Union[Notebook, NotebookId]) -> str:
         notebook_id = notebook.id if isinstance(notebook, Notebook) else notebook
         return NotebookKey(
@@ -716,8 +910,41 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             instance=self.config.platform_instance,
         ).as_urn()
 
+    def _create_upstream_class(
+        self,
+        dataset_urn: str,
+        lineage_type: Union[str, DatasetLineageTypeClass],
+        timestamp: Optional[int],
+    ) -> UpstreamClass:
+        """
+        Helper method to create UpstreamClass with optional audit stamp.
+        If timestamp is None, audit stamp is omitted.
+        """
+        if timestamp is not None:
+            return UpstreamClass(
+                dataset=dataset_urn,
+                type=lineage_type,
+                auditStamp=AuditStampClass(
+                    time=timestamp,
+                    actor=UNKNOWN_USER,
+                ),
+            )
+        else:
+            return UpstreamClass(
+                dataset=dataset_urn,
+                type=lineage_type,
+            )
+
     def gen_schema_containers(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
         domain_urn = self._gen_domain_urn(f"{schema.catalog.name}.{schema.name}")
+        schema_tags = []
+        if self.config.include_tags:
+            schema_tags = self.unity_catalog_api_proxy.get_schema_tags(
+                schema.catalog.name
+            ).get(f"{schema.catalog.name}.{schema.name}", [])
+            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 +956,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().name for tag in schema_tags]
+            if schema_tags
+            else None,
         )
 
     def gen_metastore_containers(
@@ -749,6 +979,14 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
 
     def gen_catalog_containers(self, catalog: Catalog) -> Iterable[MetadataWorkUnit]:
         domain_urn = self._gen_domain_urn(catalog.name)
+        catalog_tags = []
+        if self.config.include_tags:
+            catalog_tags = self.unity_catalog_api_proxy.get_catalog_tags(
+                catalog.name
+            ).get(catalog.name, [])
+            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 +1002,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().name for tag in catalog_tags]
+            if catalog_tags
+            else None,
         )
 
     def gen_schema_key(self, schema: Schema) -> ContainerKey:
@@ -832,6 +1073,50 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             dataset_urn=dataset_urn,
         )
 
+    def add_model_to_schema_container(
+        self, model_urn: str, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        schema_container_key = self.gen_schema_key(schema)
+        yield from add_entity_to_container(
+            container_key=schema_container_key,
+            entity_type=MlModelGroupUrn.ENTITY_TYPE,
+            entity_urn=model_urn,
+        )
+
+    def add_model_version_to_schema_container(
+        self, model_version_urn: str, schema: Schema
+    ) -> Iterable[MetadataWorkUnit]:
+        schema_container_key = self.gen_schema_key(schema)
+        yield from add_entity_to_container(
+            container_key=schema_container_key,
+            entity_type=MlModelUrn.ENTITY_TYPE,
+            entity_urn=model_version_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:
@@ -860,16 +1145,20 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         created: Optional[TimeStampClass] = None
         if table.created_at:
             custom_properties["created_at"] = str(table.created_at)
-            created = TimeStampClass(
-                int(table.created_at.timestamp() * 1000),
-                make_user_urn(table.created_by) if table.created_by else None,
-            )
+            created_ts = make_ts_millis(table.created_at)
+            if created_ts is not None:
+                created = TimeStampClass(
+                    created_ts,
+                    make_user_urn(table.created_by) if table.created_by else None,
+                )
         last_modified = created
         if table.updated_at:
-            last_modified = TimeStampClass(
-                int(table.updated_at.timestamp() * 1000),
-                table.updated_by and make_user_urn(table.updated_by),
-            )
+            updated_ts = make_ts_millis(table.updated_at)
+            if updated_ts is not None:
+                last_modified = TimeStampClass(
+                    updated_ts,
+                    table.updated_by and make_user_urn(table.updated_by),
+                )
 
         return DatasetPropertiesClass(
             name=table.name,
@@ -921,30 +1210,127 @@ 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 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:
+                    # 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"Retrieved/created platform resource for tag {tag.key.raw_text}"
+                    )
+                    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,
+                            )
+                except Exception as e:
+                    logger.exception(
+                        f"Error processing platform resource for tag {tag}"
+                    )
+                    self.report.report_warning(
+                        message="Error processing platform resource for tag",
+                        context=str(tag),
+                        title="Error processing platform resource for tag",
+                        exc=e,
+                    )
+                    continue
 
+    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:
+                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,78 +1340,26 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                     nativeDataType=column.type_text,
                     nullable=column.nullable,
                     description=column.comment,
+                    globalTags=global_tags if tags else None,
                 )
             ]
 
-    def _run_sql_parser(
-        self, view_ref: TableReference, query: str, schema_resolver: SchemaResolver
-    ) -> Optional[SqlParsingResult]:
-        raw_lineage = sqlglot_lineage(
-            query,
-            schema_resolver=schema_resolver,
-            default_db=view_ref.catalog,
-            default_schema=view_ref.schema,
-        )
-        view_urn = self.gen_dataset_urn(view_ref)
-
-        if raw_lineage.debug_info.table_error:
-            logger.debug(
-                f"Failed to parse lineage for view {view_ref}: "
-                f"{raw_lineage.debug_info.table_error}"
-            )
-            self.report.num_view_definitions_failed_parsing += 1
-            self.report.view_definitions_parsing_failures.append(
-                f"Table-level sql parsing error for view {view_ref}: {raw_lineage.debug_info.table_error}"
-            )
-            return None
-
-        elif raw_lineage.debug_info.column_error:
-            self.report.num_view_definitions_failed_column_parsing += 1
-            self.report.view_definitions_parsing_failures.append(
-                f"Column-level sql parsing error for view {view_ref}: {raw_lineage.debug_info.column_error}"
-            )
-        else:
-            self.report.num_view_definitions_parsed += 1
-            if raw_lineage.out_tables != [view_urn]:
-                self.report.num_view_definitions_view_urn_mismatch += 1
-        return view_definition_lineage_helper(raw_lineage, view_urn)
-
     def get_view_lineage(self) -> Iterable[MetadataWorkUnit]:
         if not (
             self.config.include_hive_metastore
             and self.config.include_table_lineage
-            and self.sql_parser_schema_resolver
+            and self.sql_parsing_aggregator
         ):
             return
-        # This is only used for parsing view lineage. Usage, Operations are emitted elsewhere
-        builder = SqlParsingBuilder(
-            generate_lineage=True,
-            generate_usage_statistics=False,
-            generate_operations=False,
-        )
-        for dataset_name in self.view_definitions.keys():
-            view_ref, view_definition = self.view_definitions[dataset_name]
-            result = self._run_sql_parser(
-                view_ref,
-                view_definition,
-                self.sql_parser_schema_resolver,
-            )
-            if result and result.out_tables:
-                # This does not yield any workunits but we use
-                # yield here to execute this method
-                yield from builder.process_sql_parsing_result(
-                    result=result,
-                    query=view_definition,
-                    is_view_ddl=True,
-                    include_column_lineage=self.config.include_view_column_lineage,
-                )
-        yield from builder.gen_workunits()
+
+        for mcp in self.sql_parsing_aggregator.gen_metadata():
+            yield mcp.as_workunit()
 
     def close(self):
         if self.hive_metastore_proxy:
             self.hive_metastore_proxy.close()
-        if self.view_definitions:
-            self.view_definitions.close()
+        if self.sql_parsing_aggregator:
+            self.sql_parsing_aggregator.close()
         if self.sql_parser_schema_resolver:
             self.sql_parser_schema_resolver.close()