import dataclasses
import functools
import json
import logging
import re
import time
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union, cast
from urllib.parse import urljoin

import sqlglot
import sqlglot.errors
import sqlglot.expressions as sqlglot_exp
import yaml

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
from datahub.emitter.mcp_builder import (
    CatalogKey,
    CatalogKeyWithMetastore,
    ContainerKey,
    MetastoreKey,
    NotebookKey,
    UnitySchemaKey,
    UnitySchemaKeyWithMetastore,
    add_dataset_to_container,
    add_entity_to_container,
    gen_containers,
)
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import (
    SupportStatus,
    capability,
    config_class,
    platform_name,
    support_status,
)
from datahub.ingestion.api.incremental_ownership_helper import (
    auto_incremental_ownership,
)
from datahub.ingestion.api.incremental_properties_helper import (
    auto_incremental_properties,
)
from datahub.ingestion.api.source import (
    MetadataWorkUnitProcessor,
    SourceCapability,
    TestableSource,
    TestConnectionReport,
)
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.source.aws import s3_util
from datahub.ingestion.source.aws.s3_util import (
    make_s3_urn_for_lineage,
    strip_s3_prefix,
)
from datahub.ingestion.source.common.subtypes import (
    DatasetContainerSubTypes,
    DatasetSubTypes,
    SourceCapabilityModifier,
)
from datahub.ingestion.source.state.stale_entity_removal_handler import (
    StaleEntityRemovalHandler,
)
from datahub.ingestion.source.state.stateful_ingestion_base import (
    StatefulIngestionSourceBase,
)
from datahub.ingestion.source.unity import proxy_types as unity_proxy_types
from datahub.ingestion.source.unity.analyze_profiler import UnityCatalogAnalyzeProfiler
from datahub.ingestion.source.unity.config import (
    UnityCatalogAnalyzeProfilerConfig,
    UnityCatalogGEProfilerConfig,
    UnityCatalogSourceConfig,
    UnityCatalogSQLAlchemyProfilerConfig,
)
from datahub.ingestion.source.unity.connection import create_workspace_client
from datahub.ingestion.source.unity.connection_test import UnityCatalogConnectionTest
from datahub.ingestion.source.unity.ge_profiler import UnityCatalogGEProfiler
from datahub.ingestion.source.unity.hive_metastore_proxy import (
    HIVE_METASTORE,
    HiveMetastoreProxy,
)
from datahub.ingestion.source.unity.identifier_helper import split_databricks_identifier
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,
    ServicePrincipal,
    Table,
    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 (
    GlobalTags,
    MetadataAttribution,
    Siblings,
    TagAssociation,
)
from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
    DatasetLineageType,
    FineGrainedLineage,
    FineGrainedLineageDownstreamType,
    FineGrainedLineageUpstreamType,
    Upstream,
    UpstreamLineage,
    ViewProperties,
)
from datahub.metadata.schema_classes import (
    AuditStampClass,
    BrowsePathsClass,
    DataPlatformInstanceClass,
    DatasetLineageTypeClass,
    DatasetPropertiesClass,
    DomainsClass,
    MLModelPropertiesClass,
    MySqlDDLClass,
    NullTypeClass,
    OwnerClass,
    OwnershipClass,
    OwnershipTypeClass,
    SchemaFieldClass,
    SchemaFieldDataTypeClass,
    SchemaMetadataClass,
    SubTypesClass,
    TimeStampClass,
    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.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

logger: logging.Logger = logging.getLogger(__name__)


def _parse_metric_view_source(
    raw: str,
    *,
    default_catalog: str,
    default_metastore: Optional[str],
) -> Optional[TableReference]:
    # Returns None for subqueries, 1-part names, or malformed input so the caller falls back to REST lineage.
    if not raw:
        logger.debug("metric view source rejected: empty input")
        return None
    if any(ch in raw for ch in ("(", "\n")):
        logger.debug("metric view source rejected (subquery / newline): %r", raw)
        return None
    if any(ch in raw for ch in (" ", "\t")) and "`" not in raw:
        logger.debug("metric view source rejected (unquoted whitespace): %r", raw)
        return None
    parts = split_databricks_identifier(raw)
    if parts is None:
        logger.debug("metric view source rejected (unterminated backtick): %r", raw)
        return None
    if len(parts) == 3:
        catalog, schema, name = parts
    elif len(parts) == 2:
        catalog = default_catalog
        schema, name = parts
    else:
        logger.debug(
            "metric view source rejected (need 2 or 3 parts, got %d): %r",
            len(parts),
            raw,
        )
        return None
    if not (catalog and schema and name):
        logger.debug("metric view source rejected (empty part after split): %r", raw)
        return None
    return TableReference(
        metastore=default_metastore,
        catalog=catalog,
        schema=schema,
        table=name,
    )


@platform_name("Databricks")
@config_class(UnityCatalogSourceConfig)
@capability(SourceCapability.SCHEMA_METADATA, "Enabled by default")
@capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
@capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
@capability(SourceCapability.LINEAGE_FINE, "Enabled by default")
@capability(SourceCapability.USAGE_STATS, "Enabled by default")
@capability(
    SourceCapability.OPERATION_CAPTURE,
    "Enabled by default via usage extraction, can be disabled via `include_operational_stats`",
)
@capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
@capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
@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,
    "Enabled by default via stateful ingestion",
    supported=True,
)
@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
    unity_catalog_api_proxy: UnityCatalogApiProxy
    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

    def __init__(self, ctx: PipelineContext, config: UnityCatalogSourceConfig):
        super().__init__(config, ctx)

        self.config = config
        self.report: UnityCatalogReport = UnityCatalogReport()

        self.init_hive_metastore_proxy()

        self.unity_catalog_api_proxy = UnityCatalogApiProxy(
            create_workspace_client(self.config),
            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")

        # Determine the platform_instance_name
        self.platform_instance_name = self.config.platform_instance
        if self.config.include_metastore:
            self.platform_instance_name = (
                config.workspace_name
                if config.workspace_name is not None
                else config.workspace_url.split("//")[1].split(".")[0]
            )

        if self.config.domain:
            self.domain_registry = DomainRegistry(
                cached_domains=[k for k in self.config.domain], graph=self.ctx.graph
            )

        # Global map of service principal application id -> ServicePrincipal
        self.service_principals: Dict[str, ServicePrincipal] = {}
        self.groups: List[str] = []
        # Global set of table refs
        self.table_refs: Set[TableReference] = set()
        self.view_refs: Set[TableReference] = set()
        self.notebooks: FileBackedDict[Notebook] = 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",
            )

        # SDK lacks TableType.METRIC_VIEW; warn so the no-op flag isn't silent.
        if (
            self.config.include_metric_views
            and not unity_proxy_types.metric_view_supported()
        ):
            try:
                from databricks.sdk.version import __version__ as _sdk_version
            except Exception:
                _sdk_version = "unknown"
            self.report.report_warning(
                "include_metric_views has no effect with installed databricks-sdk",
                f"databricks-sdk {_sdk_version} does not expose TableType.METRIC_VIEW. "
                "Metric views will be emitted as plain Tables. Upgrade databricks-sdk "
                "to >= 0.58.0 (the release that added this enum) to enable "
                "metric-view ingestion.",
            )

        # 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
        if self.config.include_hive_metastore:
            try:
                self.hive_metastore_proxy = HiveMetastoreProxy(
                    self.config.get_sql_alchemy_url(HIVE_METASTORE),
                    self.config.options,
                    self.report,
                )
                self.report.hive_metastore_catalog_found = True

                if self.config.include_table_lineage:
                    self.sql_parser_schema_resolver = SchemaResolver(
                        platform=self.platform,
                        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(
                    logger,
                    HIVE_METASTORE,
                    f"Failed to connect to hive_metastore due to {e}",
                )
                self.report.hive_metastore_catalog_found = False

    @staticmethod
    def test_connection(config_dict: dict) -> TestConnectionReport:
        try:
            config = UnityCatalogSourceConfig.parse_obj_allow_extras(config_dict)
        except Exception as e:
            return TestConnectionReport(
                internal_failure=True,
                internal_failure_reason=f"Failed to parse config due to {e}",
            )
        return UnityCatalogConnectionTest(config).get_connection_test()

    @classmethod
    def create(cls, config_dict, ctx):
        config = UnityCatalogSourceConfig.model_validate(config_dict)
        return cls(ctx=ctx, config=config)

    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
        return [
            *super().get_workunit_processors(),
            functools.partial(
                auto_incremental_ownership, self.config.incremental_ownership
            ),
            functools.partial(
                auto_incremental_properties, self.config.incremental_properties
            ),
            StaleEntityRemovalHandler.create(
                self, self.config, self.ctx
            ).workunit_processor,
        ]

    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
        with self.report.new_stage("Ingestion Setup"):
            wait_on_warehouse = None
            if self.config.include_hive_metastore:
                with self.report.new_stage("Start warehouse"):
                    # Can take several minutes, so start now and wait later
                    wait_on_warehouse = self.unity_catalog_api_proxy.start_warehouse()
                    if wait_on_warehouse is None:
                        self.report.report_failure(
                            "initialization",
                            f"SQL warehouse {self.config.profiling.warehouse_id} not found",
                        )
                        return
                    else:
                        # wait until warehouse is started
                        wait_on_warehouse.result()

        if self.config.include_ownership:
            with self.report.new_stage("Ingest service principals"):
                self.build_service_principal_map()
                self.build_groups_map()
        if self.config.include_notebooks:
            with self.report.new_stage("Ingest notebooks"):
                yield from self.process_notebooks()

        yield from self.process_metastores()

        yield from self.get_view_lineage()

        if self.config.include_notebooks:
            with self.report.new_stage("Notebook lineage"):
                for notebook in self.notebooks.values():
                    wu = self._gen_notebook_lineage(notebook)
                    if wu:
                        yield wu

        if self.config.include_usage_statistics:
            with self.report.new_stage("Ingest usage"):
                usage_extractor = UnityCatalogUsageExtractor(
                    config=self.config,
                    report=self.report,
                    proxy=self.unity_catalog_api_proxy,
                    table_urn_builder=self.gen_dataset_urn,
                    user_urn_builder=self.gen_user_urn,
                )
                yield from usage_extractor.get_usage_workunits(
                    self.table_refs | self.view_refs
                )

        if self.config.is_profiling_enabled():
            with self.report.new_stage("Start warehouse"):
                # Need to start the warehouse again for profiling,
                # as it may have been stopped after ingestion might take
                # longer time to complete
                wait_on_warehouse = self.unity_catalog_api_proxy.start_warehouse()
                if wait_on_warehouse is None:
                    self.report.report_failure(
                        "initialization",
                        f"SQL warehouse {self.config.profiling.warehouse_id} not found",
                    )
                    return
                else:
                    # wait until warehouse is started
                    wait_on_warehouse.result()

            with self.report.new_stage("Profiling"):
                if isinstance(self.config.profiling, UnityCatalogAnalyzeProfilerConfig):
                    logger.info(
                        "Using UnityCatalogAnalyzeProfiler (ANALYZE TABLE method)"
                    )
                    yield from UnityCatalogAnalyzeProfiler(
                        self.config.profiling,
                        self.report,
                        self.unity_catalog_api_proxy,
                        self.gen_dataset_urn,
                    ).get_workunits(self.table_refs)
                elif isinstance(
                    self.config.profiling, UnityCatalogSQLAlchemyProfilerConfig
                ):
                    logger.info(
                        "Using UnityCatalogGEProfiler with SQLAlchemyProfiler (method: sqlalchemy)"
                    )
                    # Use GenericProfiler which will use SQLAlchemyProfiler internally
                    yield from UnityCatalogGEProfiler(
                        config=self.config,
                        profiling_config=self.config.profiling,
                        report=self.report,
                    ).get_workunits(list(self.tables.values()))
                elif isinstance(self.config.profiling, UnityCatalogGEProfilerConfig):
                    logger.info(
                        "Using UnityCatalogGEProfiler with DatahubGEProfiler (method: ge)"
                    )
                    yield from UnityCatalogGEProfiler(
                        config=self.config,
                        profiling_config=self.config.profiling,
                        report=self.report,
                    ).get_workunits(list(self.tables.values()))
                else:
                    raise ValueError("Unknown profiling config method")

    def build_service_principal_map(self) -> None:
        try:
            for sp in self.unity_catalog_api_proxy.service_principals():
                self.service_principals[sp.application_id] = sp
        except Exception as e:
            self.report.report_warning(
                "service-principals", f"Unable to fetch service principals: {e}"
            )

    def build_groups_map(self) -> None:
        try:
            self.groups += self.unity_catalog_api_proxy.groups()
        except Exception as e:
            self.report.report_warning("groups", f"Unable to fetch groups: {e}")

    def process_notebooks(self) -> Iterable[MetadataWorkUnit]:
        for notebook in self.unity_catalog_api_proxy.workspace_notebooks():
            if not self.config.notebook_pattern.allowed(notebook.path):
                self.report.notebooks.dropped(notebook.path)
                continue

            self.notebooks[str(notebook.id)] = notebook
            yield from self._gen_notebook_workunits(notebook)

    def _gen_notebook_workunits(self, notebook: Notebook) -> Iterable[MetadataWorkUnit]:
        properties = {"path": notebook.path}
        if notebook.language:
            properties["language"] = notebook.language.value

        mcps = MetadataChangeProposalWrapper.construct_many(
            entityUrn=self.gen_notebook_urn(notebook),
            aspects=[
                DatasetPropertiesClass(
                    name=notebook.path.rsplit("/", 1)[-1],
                    customProperties=properties,
                    externalUrl=urljoin(
                        self.config.workspace_url, f"#notebook/{notebook.id}"
                    ),
                    created=(
                        TimeStampClass(make_ts_millis(notebook.created_at))
                        if notebook.created_at
                        else None
                    ),
                    lastModified=(
                        TimeStampClass(make_ts_millis(notebook.modified_at))
                        if notebook.modified_at
                        else None
                    ),
                ),
                SubTypesClass(typeNames=[DatasetSubTypes.NOTEBOOK]),
                BrowsePathsClass(paths=notebook.path.split("/")),
                self._create_data_platform_instance_aspect(),
            ],
        )
        for mcp in mcps:
            yield mcp.as_workunit()

        self.report.notebooks.processed(notebook.path)

    def _gen_notebook_lineage(self, notebook: Notebook) -> Optional[MetadataWorkUnit]:
        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=upstreams),
        ).as_workunit()

    def process_metastores(self) -> Iterable[MetadataWorkUnit]:
        metastore: Optional[Metastore] = None
        if self.config.include_metastore:
            metastore = self.unity_catalog_api_proxy.assigned_metastore()
            if not metastore:
                self.report.report_failure("Metastore", "Not found")
                return
            yield from self.gen_metastore_containers(metastore)
        yield from self.process_catalogs(metastore)
        if metastore and self.config.include_metastore:
            self.report.metastores.processed(metastore.id)

    def process_catalogs(
        self, metastore: Optional[Metastore]
    ) -> Iterable[MetadataWorkUnit]:
        for catalog in self._get_catalogs(metastore):
            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)

                self.report.catalogs.processed(catalog.id)

    def _get_catalogs(self, metastore: Optional[Metastore]) -> Iterable[Catalog]:
        if self.config.catalogs:
            for catalog_name in self.config.catalogs:
                catalog = self.unity_catalog_api_proxy.catalog(
                    catalog_name, metastore=metastore
                )
                if catalog:
                    yield catalog
        else:
            yield from self.unity_catalog_api_proxy.catalogs(metastore=metastore)

    def process_schemas(self, catalog: Catalog) -> Iterable[MetadataWorkUnit]:
        for schema in self.unity_catalog_api_proxy.schemas(catalog=catalog):
            if not self.config.schema_pattern.allowed(schema.id):
                self.report.schemas.dropped(schema.id)
                continue

            with self.report.new_stage(f"Ingest schema {schema.id}"):
                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(
                        message="Missed schema because of parsing issues",
                        context=str(schema),
                        title="Error parsing schema",
                        exc=e,
                    )
                    continue

                self.report.schemas.processed(schema.id)

    def process_tables(self, schema: Schema) -> Iterable[MetadataWorkUnit]:
        for table in self.unity_catalog_api_proxy.tables(schema=schema):
            if not self.config.table_pattern.allowed(table.ref.qualified_table_name):
                self.report.tables.dropped(table.id, f"table ({table.table_type})")
                continue

            if (
                table.is_metric_view
                and self.config.include_metric_views
                and not self.config.metric_view_pattern.allowed(
                    table.ref.qualified_table_name
                )
            ):
                self.report.tables.dropped(table.id, f"table ({table.table_type})")
                self.report.metric_views.dropped(table.id)
                continue

            if (
                self.config.is_profiling_enabled()
                and self.config.is_ge_profiling()
                and self.config.profiling.pattern.allowed(
                    table.ref.qualified_table_name
                )
                and not table.is_view
            ):
                self.tables[table.ref.qualified_table_name] = table

            if table.is_view:
                self.view_refs.add(table.ref)
            else:
                # Keep metric views in table_refs so stateful soft-delete still tracks them.
                self.table_refs.add(table.ref)
            yield from self.process_table(table, schema)
            self.report.tables.processed(table.id, f"table ({table.table_type})")
            if table.is_metric_view and self.config.include_metric_views:
                self.report.metric_views.processed(table.id)

    def process_table(self, table: Table, schema: Schema) -> Iterable[MetadataWorkUnit]:
        dataset_urn = self.gen_dataset_urn(table.ref)
        yield from self.add_table_to_dataset_container(dataset_urn, schema)

        metric_view_spec: Optional[Dict[str, Any]] = None
        if table.is_metric_view and self.config.include_metric_views:
            metric_view_spec = self._load_metric_view_spec(table)

        table_props = self._create_table_property_aspect(table, metric_view_spec)
        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, metric_view_spec)

        sub_type = self._create_table_sub_type_aspect(table)
        schema_metadata, platform_resources = self._create_schema_metadata_aspect(
            table, metric_view_spec
        )
        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()

        if table.is_metric_view and self.config.include_metric_views:
            # REST fallback so subquery sources / parse failures don't lose upstreams.
            lineage = (
                self._extract_metric_view_lineage(table, metric_view_spec)
                if metric_view_spec is not None
                else None
            )
            if lineage is None:
                lineage = self.ingest_lineage(table)
                if lineage is None:
                    self.report.report_warning(
                        "Metric view has no upstream lineage",
                        f"{table.ref.qualified_table_name}: neither YAML parsing "
                        "nor the table-lineage REST API produced any upstreams.",
                    )
        else:
            lineage = self.ingest_lineage(table)

        if self.config.include_notebooks:
            for notebook_id in table.downstream_notebooks:
                if str(notebook_id) in self.notebooks:
                    self.notebooks[str(notebook_id)] = Notebook.add_upstream(
                        table.ref, self.notebooks[str(notebook_id)]
                    )

        # Sql parsing is required only for hive metastore view lineage
        if (
            self.sql_parser_schema_resolver
            and table.schema.catalog.type == CustomCatalogType.HIVE_METASTORE_CATALOG
        ):
            self.sql_parser_schema_resolver.add_schema_metadata(
                dataset_urn, schema_metadata
            )
            if (
                table.view_definition
                and self.sql_parsing_aggregator
                and not (table.is_metric_view and self.config.include_metric_views)
            ):
                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")
            in {"EXTERNAL", "HIVE_EXTERNAL_TABLE"}
            and table_props.customProperties.get("data_source_format") == "DELTA"
            and self.config.emit_siblings
        ):
            storage_location = str(table_props.customProperties.get("storage_location"))
            if any(
                storage_location.startswith(prefix) for prefix in s3_util.S3_PREFIXES
            ):
                browse_path = strip_s3_prefix(storage_location)
                source_dataset_urn = make_dataset_urn_with_platform_instance(
                    "delta-lake",
                    browse_path,
                    self.config.delta_lake_options.platform_instance_name,
                    self.config.delta_lake_options.env,
                )

                yield from self.gen_siblings_workunit(dataset_urn, source_dataset_urn)
                yield from self.gen_lineage_workunit(dataset_urn, source_dataset_urn)

        yield from [
            mcp.as_workunit()
            for mcp in MetadataChangeProposalWrapper.construct_many(
                entityUrn=dataset_urn,
                aspects=[
                    table_props,
                    view_props,
                    sub_type,
                    schema_metadata,
                    domain,
                    ownership,
                    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),
                )
            )
        custom_properties = {}
        if ml_model_version.signature:
            for key, value in dataclasses.asdict(ml_model_version.signature).items():
                if value:
                    custom_properties[f"signature.{key}"] = json.dumps(value)

        if ml_model_version.run_details:
            if ml_model_version.run_details.tags:
                for key, value in ml_model_version.run_details.tags.items():
                    if value:
                        custom_properties[key] = json.dumps(value)

        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,
            training_metrics=cast(
                Optional[Dict[str, Optional[str]]], ml_model_version.run_details.metrics
            )
            if ml_model_version.run_details and ml_model_version.run_details.metrics
            else None,
            hyper_params=cast(
                Optional[Dict[str, Optional[str]]],
                ml_model_version.run_details.parameters,
            )
            if ml_model_version.run_details and ml_model_version.run_details.parameters
            else None,
            custom_properties=custom_properties if custom_properties else None,
            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,
                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

            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)

    def _generate_lineage_aspect(
        self, dataset_urn: str, table: Table
    ) -> Optional[UpstreamLineageClass]:
        upstreams: List[UpstreamClass] = []
        finegrained_lineages: List[FineGrainedLineage] = []
        for upstream_ref, downstream_to_upstream_cols in sorted(
            table.upstreams.items()
        ):
            upstream_urn = self.gen_dataset_urn(upstream_ref)

            # Should be empty if config.include_column_lineage is False
            finegrained_lineages.extend(
                FineGrainedLineage(
                    upstreamType=FineGrainedLineageUpstreamType.FIELD_SET,
                    upstreams=[
                        make_schema_field_urn(upstream_urn, upstream_col)
                        for upstream_col in sorted(u_cols)
                    ],
                    downstreamType=FineGrainedLineageUpstreamType.FIELD_SET,
                    downstreams=[make_schema_field_urn(dataset_urn, d_col)],
                )
                for d_col, u_cols in sorted(downstream_to_upstream_cols.items())
            )

            timestamp = make_ts_millis(upstream_ref.last_updated)
            upstreams.append(
                self._create_upstream_class(
                    upstream_urn,
                    DatasetLineageTypeClass.TRANSFORMED,
                    timestamp,
                )
            )

        for notebook in table.upstream_notebooks.values():
            timestamp = make_ts_millis(notebook.last_updated)
            upstreams.append(
                self._create_upstream_class(
                    self.gen_notebook_urn(notebook.id),
                    DatasetLineageTypeClass.TRANSFORMED,
                    timestamp,
                )
            )

        if self.config.include_external_lineage:
            for external_ref in table.external_upstreams:
                if not external_ref.has_permission or not external_ref.path:
                    self.report.num_external_upstreams_lacking_permissions += 1
                    logger.warning(
                        f"Lacking permissions for external file upstream on {table.ref}"
                    )
                elif external_ref.path.startswith("s3://"):
                    upstreams.append(
                        UpstreamClass(
                            dataset=make_s3_urn_for_lineage(
                                external_ref.path, self.config.env
                            ),
                            type=DatasetLineageTypeClass.COPY,
                        )
                    )
                else:
                    self.report.num_external_upstreams_unsupported += 1
                    logger.warning(
                        f"Unsupported external file upstream on {table.ref}: {external_ref.path}"
                    )

        if upstreams:
            return UpstreamLineageClass(
                upstreams=upstreams,
                fineGrainedLineages=(
                    finegrained_lineages if self.config.include_column_lineage else None
                ),
            )
        else:
            return None

    def _get_domain_aspect(self, dataset_name: str) -> Optional[DomainsClass]:
        domain_urn = self._gen_domain_urn(dataset_name)
        if not domain_urn:
            return None
        return DomainsClass(domains=[domain_urn])

    def get_owner_urn(self, user: Optional[str]) -> Optional[str]:
        if self.config.include_ownership and user is not None:
            if user in self.groups:
                return make_group_urn(user)
            return self.gen_user_urn(user)
        return None

    def gen_user_urn(self, user: str) -> str:
        if user in self.service_principals:
            user = self.service_principals[user].display_name
        return make_user_urn(user)

    def gen_dataset_urn(self, table_ref: TableReference) -> str:
        return make_dataset_urn_with_platform_instance(
            platform=self.platform,
            platform_instance=self.platform_instance_name,
            name=str(table_ref),
            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(
            notebook_id=notebook_id,
            platform=self.platform,
            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(
            container_key=schema_container_key,
            name=schema.name,
            sub_types=[DatasetContainerSubTypes.SCHEMA],
            parent_container_key=self.gen_catalog_key(catalog=schema.catalog),
            domain_urn=domain_urn,
            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(
        self, metastore: Metastore
    ) -> Iterable[MetadataWorkUnit]:
        domain_urn = self._gen_domain_urn(metastore.name)

        metastore_container_key = self.gen_metastore_key(metastore)
        yield from gen_containers(
            container_key=metastore_container_key,
            name=metastore.name,
            sub_types=[DatasetContainerSubTypes.DATABRICKS_METASTORE],
            domain_urn=domain_urn,
            description=metastore.comment,
            owner_urn=self.get_owner_urn(metastore.owner),
            external_url=self.external_url_base,
        )

    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(
            container_key=catalog_container_key,
            name=catalog.name,
            sub_types=[DatasetContainerSubTypes.CATALOG],
            domain_urn=domain_urn,
            parent_container_key=(
                self.gen_metastore_key(catalog.metastore)
                if self.config.include_metastore and catalog.metastore
                else None
            ),
            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:
        if self.config.include_metastore:
            assert schema.catalog.metastore
            return UnitySchemaKeyWithMetastore(
                unity_schema=schema.name,
                platform=self.platform,
                instance=self.config.platform_instance,
                catalog=schema.catalog.name,
                metastore=schema.catalog.metastore.name,
                env=self.config.env,
            )
        else:
            return UnitySchemaKey(
                unity_schema=schema.name,
                platform=self.platform,
                instance=self.config.platform_instance,
                catalog=schema.catalog.name,
                env=self.config.env,
            )

    def gen_metastore_key(self, metastore: Metastore) -> MetastoreKey:
        return MetastoreKey(
            metastore=metastore.name,
            platform=self.platform,
            instance=self.config.platform_instance,
            env=self.config.env,
        )

    def gen_catalog_key(self, catalog: Catalog) -> ContainerKey:
        if self.config.include_metastore:
            assert catalog.metastore
            return CatalogKeyWithMetastore(
                catalog=catalog.name,
                metastore=catalog.metastore.name,
                platform=self.platform,
                instance=self.config.platform_instance,
                env=self.config.env,
            )
        else:
            return CatalogKey(
                catalog=catalog.name,
                platform=self.platform,
                instance=self.config.platform_instance,
                env=self.config.env,
            )

    def _gen_domain_urn(self, dataset_name: str) -> Optional[str]:
        domain_urn: Optional[str] = None

        for domain, pattern in self.config.domain.items():
            if pattern.allowed(dataset_name):
                domain_urn = make_domain_urn(
                    self.domain_registry.get_domain_urn(domain)
                )

        return domain_urn

    def add_table_to_dataset_container(
        self, dataset_urn: str, schema: Schema
    ) -> Iterable[MetadataWorkUnit]:
        schema_container_key = self.gen_schema_key(schema)
        yield from add_dataset_to_container(
            container_key=schema_container_key,
            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,
        metric_view_spec: Optional[Dict[str, Any]] = None,
    ) -> DatasetPropertiesClass:
        custom_properties: dict = {}
        if table.storage_location is not None:
            custom_properties["storage_location"] = table.storage_location
        if table.data_source_format is not None:
            custom_properties["data_source_format"] = table.data_source_format.value
        if table.generation is not None:
            custom_properties["generation"] = str(table.generation)

        if table.table_type:
            custom_properties["table_type"] = table.table_type.value

        if table.created_by:
            custom_properties["created_by"] = table.created_by
        if table.properties:
            # Unity Catalog injects ~150 view.sqlConfig.spark.* keys per metric view; drop them.
            drop_sql_config = table.is_metric_view and self.config.include_metric_views
            for k, v in table.properties.items():
                if drop_sql_config and k.startswith("view.sqlConfig.spark."):
                    continue
                custom_properties[k] = str(v)
        if table.table_id:
            custom_properties["table_id"] = table.table_id
        if table.owner:
            custom_properties["owner"] = table.owner
        if table.updated_by:
            custom_properties["updated_by"] = table.updated_by
        if table.updated_at:
            custom_properties["updated_at"] = str(table.updated_at)

        if (
            table.is_metric_view
            and self.config.include_metric_views
            and metric_view_spec is not None
        ):
            filter_expr = metric_view_spec.get("filter")
            if isinstance(filter_expr, str) and filter_expr.strip():
                custom_properties["metric_view.filter"] = filter_expr

        created: Optional[TimeStampClass] = None
        if table.created_at:
            custom_properties["created_at"] = str(table.created_at)
            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:
            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,
            qualifiedName=table.ref.qualified_table_name,
            description=table.comment,
            customProperties=custom_properties,
            created=created,
            lastModified=last_modified,
            externalUrl=f"{self.external_url_base}/{table.ref.external_path}",
        )

    def _create_table_ownership_aspect(self, table: Table) -> Optional[OwnershipClass]:
        owner_urn = self.get_owner_urn(table.owner)
        if owner_urn is not None:
            return OwnershipClass(
                owners=[
                    OwnerClass(
                        owner=owner_urn,
                        type=OwnershipTypeClass.DATAOWNER,
                    )
                ]
            )
        return None

    def _create_data_platform_instance_aspect(
        self,
    ) -> Optional[DataPlatformInstanceClass]:
        if self.config.ingest_data_platform_instance_aspect:
            return DataPlatformInstanceClass(
                platform=make_data_platform_urn(self.platform),
                instance=(
                    make_dataplatform_instance_urn(
                        self.platform, self.platform_instance_name
                    )
                    if self.platform_instance_name
                    else None
                ),
            )
        return None

    def _create_table_sub_type_aspect(self, table: Table) -> SubTypesClass:
        if table.is_metric_view and self.config.include_metric_views:
            type_name = DatasetSubTypes.METRIC_VIEW
        elif table.is_view:
            type_name = DatasetSubTypes.VIEW
        else:
            type_name = DatasetSubTypes.TABLE
        return SubTypesClass(typeNames=[type_name])

    def _create_view_property_aspect(
        self, table: Table, metric_view_spec: Optional[Dict[str, Any]] = None
    ) -> ViewProperties:
        assert table.view_definition
        materialized = False
        if table.is_metric_view and self.config.include_metric_views:
            view_language = "YAML"
            if metric_view_spec is not None:
                mat = metric_view_spec.get("materialization")
                materialized = isinstance(mat, str) and mat.lower() == "materialized"
        else:
            view_language = "SQL"
        return ViewProperties(
            materialized=materialized,
            viewLanguage=view_language,
            viewLogic=table.view_definition,
        )

    def _load_metric_view_spec(self, table: Table) -> Optional[Dict[str, Any]]:
        consequence = (
            "Falling back to the REST table-lineage API; column-level lineage "
            "and Dimension/Measure tags will be absent for this view."
        )
        if not table.view_definition:
            self.report.num_metric_views_yaml_shape_invalid += 1
            self.report.report_warning(
                "Metric view has no YAML body",
                f"{table.ref.qualified_table_name}: view_definition is empty. "
                f"{consequence}",
            )
            return None
        try:
            spec = yaml.safe_load(table.view_definition)
        except (yaml.YAMLError, RecursionError, UnicodeDecodeError) as e:
            self.report.num_metric_views_yaml_parse_failures += 1
            self.report.report_warning(
                "Metric view YAML failed to parse",
                f"{table.ref.qualified_table_name}: {type(e).__name__}: {e}. "
                f"{consequence}",
            )
            return None
        if not isinstance(spec, dict):
            self.report.num_metric_views_yaml_shape_invalid += 1
            self.report.report_warning(
                "Metric view YAML is not a mapping",
                f"{table.ref.qualified_table_name}: top-level type is "
                f"{type(spec).__name__!r}, expected dict. {consequence}",
            )
            return None
        return spec

    def _extract_metric_view_lineage(
        self,
        table: Table,
        metric_view_spec: Optional[Dict[str, Any]] = None,
    ) -> Optional[UpstreamLineageClass]:
        if metric_view_spec is None:
            metric_view_spec = self._load_metric_view_spec(table)
        if metric_view_spec is None:
            return None

        shape_issues: List[str] = []
        source_names: List[str] = []
        top_source = metric_view_spec.get("source")
        if isinstance(top_source, str):
            source_names.append(top_source)
        elif top_source is not None:
            shape_issues.append(
                f"`source` is {type(top_source).__name__}, expected str"
            )

        joins = metric_view_spec.get("joins")
        if isinstance(joins, list):
            for i, join in enumerate(joins):
                if not isinstance(join, dict):
                    shape_issues.append(
                        f"`joins[{i}]` is {type(join).__name__}, expected mapping"
                    )
                    continue
                js = join.get("source")
                if isinstance(js, str):
                    source_names.append(js)
                elif js is not None:
                    shape_issues.append(
                        f"`joins[{i}].source` is {type(js).__name__}, expected str"
                    )
        elif joins is not None:
            shape_issues.append(f"`joins` is {type(joins).__name__}, expected list")

        if shape_issues:
            self.report.num_metric_views_yaml_shape_invalid += 1
            self.report.report_warning(
                "Metric view YAML has unexpected shape",
                f"{table.ref.qualified_table_name}: {'; '.join(shape_issues)}. "
                "Affected entries are skipped; remaining YAML continues to be processed.",
            )

        default_catalog = table.ref.catalog
        default_metastore = table.ref.metastore
        upstreams: List[UpstreamClass] = []
        unparseable: List[str] = []
        for raw in source_names:
            ref = _parse_metric_view_source(
                raw,
                default_catalog=default_catalog,
                default_metastore=default_metastore,
            )
            if ref is None:
                unparseable.append(raw)
                continue
            upstreams.append(
                self._create_upstream_class(
                    self.gen_dataset_urn(ref),
                    DatasetLineageTypeClass.TRANSFORMED,
                    None,
                )
            )

        if unparseable and upstreams:
            # All-rejected case is reported below via num_metric_views_no_parseable_sources.
            self.report.num_metric_view_unparseable_sources += len(unparseable)
            self.report.report_warning(
                "Metric view source(s) skipped",
                f"{table.ref.qualified_table_name}: unparseable sources="
                f"{unparseable!r}. Need 2-part (schema.table) or 3-part "
                "(catalog.schema.table) identifiers; 1-part names and SQL "
                "subqueries are not supported by the YAML lineage path.",
            )

        if not upstreams:
            if source_names:
                self.report.num_metric_views_no_parseable_sources += 1
                self.report.report_warning(
                    "Metric view has no parseable upstream sources",
                    f"{table.ref.qualified_table_name}: sources={source_names!r}. "
                    "Lineage extraction needs 2-part (schema.table) or 3-part "
                    "(catalog.schema.table) identifiers; 1-part names and SQL "
                    "subqueries fall back to the REST table-lineage API.",
                )
            return None

        fine_grained: Optional[List[FineGrainedLineage]] = None
        if self.config.include_column_lineage:
            fine_grained = (
                self._extract_metric_view_column_lineage(table, metric_view_spec)
                or None
            )
        return UpstreamLineageClass(
            upstreams=upstreams,
            fineGrainedLineages=fine_grained,
        )

    def _extract_metric_view_column_lineage(
        self, table: Table, spec: Dict[str, Any]
    ) -> List[FineGrainedLineage]:
        default_catalog = table.ref.catalog
        default_metastore = table.ref.metastore
        alias_map: Dict[str, TableReference] = {}
        source_raw = spec.get("source")
        if isinstance(source_raw, str):
            source_ref = _parse_metric_view_source(
                source_raw,
                default_catalog=default_catalog,
                default_metastore=default_metastore,
            )
            if source_ref is not None:
                alias_map[""] = source_ref

        joins_skipped = 0
        joins = spec.get("joins")
        if isinstance(joins, list):
            for join in joins:
                if not isinstance(join, dict):
                    joins_skipped += 1
                    continue
                join_name = join.get("name")
                join_source = join.get("source")
                if not isinstance(join_name, str) or not isinstance(join_source, str):
                    joins_skipped += 1
                    continue
                ref = _parse_metric_view_source(
                    join_source,
                    default_catalog=default_catalog,
                    default_metastore=default_metastore,
                )
                if ref is None:
                    joins_skipped += 1
                    continue
                alias_map[join_name.lower()] = ref

        if joins_skipped:
            self.report.num_metric_view_joins_skipped += joins_skipped
            self.report.report_warning(
                "Metric view joins skipped",
                f"{table.ref.qualified_table_name}: {joins_skipped} join "
                "entry/entries skipped due to malformed shape or unparseable "
                "source. Column lineage for join-qualified columns will be absent.",
            )

        downstream_urn = self.gen_dataset_urn(table.ref)
        fine: List[FineGrainedLineage] = []
        unresolved_qualifiers: Set[str] = set()
        skipped_entries = 0
        for key in ("dimensions", "measures"):
            entries = spec.get(key)
            if not isinstance(entries, list):
                continue
            for entry in entries:
                if not isinstance(entry, dict):
                    skipped_entries += 1
                    continue
                col_name = entry.get("name")
                expr = entry.get("expr")
                if not isinstance(col_name, str) or not isinstance(expr, str):
                    skipped_entries += 1
                    continue
                upstream_cols, unresolved = self._parse_metric_view_expr(
                    expr,
                    alias_map,
                    context=f"{table.ref.qualified_table_name}:{col_name}",
                )
                unresolved_qualifiers.update(unresolved)
                if not upstream_cols:
                    continue
                upstream_urns = [
                    make_schema_field_urn(self.gen_dataset_urn(ref), upc)
                    for (ref, upc) in upstream_cols
                ]
                fine.append(
                    FineGrainedLineage(
                        upstreamType=FineGrainedLineageUpstreamType.FIELD_SET,
                        upstreams=upstream_urns,
                        downstreamType=FineGrainedLineageDownstreamType.FIELD,
                        downstreams=[make_schema_field_urn(downstream_urn, col_name)],
                    )
                )

        if skipped_entries:
            self.report.num_metric_view_skipped_dim_measure_entries += skipped_entries
            self.report.report_warning(
                "Metric view dimension/measure entries skipped",
                f"{table.ref.qualified_table_name}: {skipped_entries} "
                "dimension/measure entry/entries skipped due to malformed shape "
                "(non-mapping entry, missing or non-string `name`/`expr`). Column "
                "lineage for affected columns will be absent.",
            )

        if unresolved_qualifiers:
            self.report.num_metric_view_unresolved_qualifiers += len(
                unresolved_qualifiers
            )
            self.report.report_warning(
                "Metric view expression references unknown qualifier",
                f"{table.ref.qualified_table_name}: qualifiers="
                f"{sorted(unresolved_qualifiers)!r} not found in `source`/`joins`. "
                "Lineage rows for affected columns were skipped.",
            )

        return fine

    def _parse_metric_view_expr(
        self,
        expr: str,
        alias_map: Dict[str, TableReference],
        context: str = "",
    ) -> Tuple[List[Tuple[TableReference, str]], Set[str]]:
        try:
            # sqlglot internals can leak AttributeError/TypeError on malformed ASTs.
            tree = sqlglot.parse_one(expr, dialect="databricks")
        except (
            sqlglot.errors.SqlglotError,
            RecursionError,
            ValueError,
            AttributeError,
            TypeError,
        ) as e:
            self.report.num_metric_views_expr_parse_failures += 1
            self.report.report_warning(
                "Metric view expression failed to parse",
                f"{context}: expr={expr!r}: {type(e).__name__}: {e}".lstrip(": "),
            )
            return [], set()
        if tree is None:
            self.report.num_metric_view_expr_empty_tree += 1
            return [], set()

        result: List[Tuple[TableReference, str]] = []
        unresolved: Set[str] = set()
        seen: Set[Tuple[str, str]] = set()
        for col in tree.find_all(sqlglot_exp.Column):
            qualifier = col.table.lower() if col.table else ""
            ref = alias_map.get(qualifier)
            if ref is None:
                # "<source>" sentinel surfaces unqualified columns when top `source`
                # was unparseable; without it they would be silently dropped.
                unresolved.add(qualifier or "<source>")
                continue
            col_name = col.name
            key = (str(ref), col_name)
            if key in seen:
                continue
            seen.add(key)
            result.append((ref, col_name))
        return result, unresolved

    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,
        metric_view_spec: Optional[Dict[str, Any]] = None,
    ) -> Tuple[SchemaMetadataClass, Iterable[MetadataWorkUnit]]:
        schema_fields: List[SchemaFieldClass] = []
        unique_tags: Set[UnityCatalogTag] = set()
        dim_measure_tags, dim_measure_descriptions = (
            self._build_metric_view_column_overrides(table, metric_view_spec)
        )
        for column in table.columns:
            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]
            extra_tag = dim_measure_tags.get(column.name.lower())
            description_override = dim_measure_descriptions.get(column.name.lower())
            schema_fields.extend(
                self._create_schema_field(
                    column,
                    tag_urns,
                    extra_tag=extra_tag,
                    description_override=description_override,
                )
            )

        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,
        )

    def _build_metric_view_column_overrides(
        self, table: Table, metric_view_spec: Optional[Dict[str, Any]]
    ) -> Tuple[Dict[str, TagUrn], Dict[str, str]]:
        if not (
            table.is_metric_view
            and self.config.include_metric_views
            and metric_view_spec is not None
        ):
            return {}, {}
        tag_map: Dict[str, TagUrn] = {}
        desc_map: Dict[str, str] = {}
        for key, tag in (("dimensions", "Dimension"), ("measures", "Measure")):
            entries = metric_view_spec.get(key)
            if not isinstance(entries, list):
                continue
            for entry in entries:
                if not isinstance(entry, dict):
                    continue
                name = entry.get("name")
                if not isinstance(name, str) or not name:
                    continue
                tag_map[name.lower()] = TagUrn(tag)
                description = entry.get("description")
                if isinstance(description, str) and description.strip():
                    desc_map[name.lower()] = description
        return tag_map, desc_map

    @staticmethod
    def _create_schema_field(
        column: Column,
        tags: Optional[List[TagUrn]],
        extra_tag: Optional[TagUrn] = None,
        description_override: Optional[str] = None,
    ) -> List[SchemaFieldClass]:
        _COMPLEX_TYPE = re.compile("^(struct|array)")
        attribution = MetadataAttribution(
            source="urn:li:dataPlatform:unity-catalog",
            actor="urn:li:corpuser:datahub",
            time=int(time.time() * 1000),
        )
        description = description_override if description_override else column.comment

        if _COMPLEX_TYPE.match(column.type_text.lower()):
            fields = get_schema_fields_for_hive_column(
                column.name, column.type_text.lower(), description=description
            )
            # Complex types: tag the root field only; children carry no semantics.
            if fields and extra_tag is not None:
                fields[0].globalTags = GlobalTags(
                    tags=[TagAssociation(tag=extra_tag.urn(), attribution=attribution)]
                )
            return fields

        tag_associations: List[TagAssociation] = []
        if tags is not None:
            tag_associations.extend(
                TagAssociation(tag=tag.urn(), attribution=attribution) for tag in tags
            )
        if extra_tag is not None:
            tag_associations.append(
                TagAssociation(tag=extra_tag.urn(), attribution=attribution)
            )
        global_tags = GlobalTags(tags=tag_associations) if tag_associations else None

        return [
            SchemaFieldClass(
                fieldPath=column.name,
                type=SchemaFieldDataTypeClass(
                    type=DATA_TYPE_REGISTRY.get(column.type_name, NullTypeClass)()
                ),
                nativeDataType=column.type_text,
                nullable=column.nullable,
                description=description,
                globalTags=global_tags,
            )
        ]

    def get_view_lineage(self) -> Iterable[MetadataWorkUnit]:
        if not (
            self.config.include_hive_metastore
            and self.config.include_table_lineage
            and self.sql_parsing_aggregator
        ):
            return

        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.sql_parsing_aggregator:
            self.sql_parsing_aggregator.close()
        if self.sql_parser_schema_resolver:
            self.sql_parser_schema_resolver.close()

        super().close()

    def gen_siblings_workunit(
        self,
        dataset_urn: str,
        source_dataset_urn: str,
    ) -> Iterable[MetadataWorkUnit]:
        """
        Generate sibling workunit for both unity-catalog dataset and its connector source dataset
        """
        yield MetadataChangeProposalWrapper(
            entityUrn=dataset_urn,
            aspect=Siblings(primary=False, siblings=[source_dataset_urn]),
        ).as_workunit()

        yield MetadataChangeProposalWrapper(
            entityUrn=source_dataset_urn,
            aspect=Siblings(primary=True, siblings=[dataset_urn]),
        ).as_workunit(is_primary_source=False)

    def gen_lineage_workunit(
        self,
        dataset_urn: str,
        source_dataset_urn: str,
    ) -> Iterable[MetadataWorkUnit]:
        """
        Generate dataset to source connector lineage workunit
        """
        yield MetadataChangeProposalWrapper(
            entityUrn=dataset_urn,
            aspect=UpstreamLineage(
                upstreams=[
                    Upstream(dataset=source_dataset_urn, type=DatasetLineageType.VIEW)
                ]
            ),
        ).as_workunit()
