acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/unity/proxy.py

@@ -4,16 +4,22 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 
 import dataclasses
 import logging
+import os
+from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
-from typing import Any, Dict, Iterable, List, Optional, Union, cast
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union, cast
 from unittest.mock import patch
 
+import cachetools
+from cachetools import cached
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service.catalog import (
     CatalogInfo,
     ColumnInfo,
     GetMetastoreSummaryResponse,
     MetastoreInfo,
+    ModelVersionInfo,
+    RegisteredModelInfo,
     SchemaInfo,
     TableInfo,
 )
@@ -25,9 +31,17 @@ from databricks.sdk.service.sql import (
     QueryStatus,
 )
 from databricks.sdk.service.workspace import ObjectType
+from databricks.sql import connect
+from databricks.sql.types import Row
+from typing_extensions import assert_never
 
 from datahub._version import nice_version_name
+from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
 from datahub.emitter.mce_builder import parse_ts_millis
+from datahub.ingestion.source.unity.config import (
+    LineageDataSource,
+    UsageDataSource,
+)
 from datahub.ingestion.source.unity.hive_metastore_proxy import HiveMetastoreProxy
 from datahub.ingestion.source.unity.proxy_profiling import (
     UnityCatalogProxyProfilingMixin,
@@ -39,7 +53,10 @@ from datahub.ingestion.source.unity.proxy_types import (
     CustomCatalogType,
     ExternalTableReference,
     Metastore,
+    Model,
+    ModelVersion,
     Notebook,
+    NotebookReference,
     Query,
     Schema,
     ServicePrincipal,
@@ -47,9 +64,31 @@ from datahub.ingestion.source.unity.proxy_types import (
     TableReference,
 )
 from datahub.ingestion.source.unity.report import UnityCatalogReport
+from datahub.utilities.file_backed_collections import FileBackedDict
 
 logger: logging.Logger = logging.getLogger(__name__)
 
+# It is enough to keep the cache size to 1, since we only process one catalog at a time
+# We need to change this if we want to support parallel processing of multiple catalogs
+_MAX_CONCURRENT_CATALOGS = 1
+
+
+# Import and apply the proxy patch from separate module
+try:
+    from datahub.ingestion.source.unity.proxy_patch import (
+        apply_databricks_proxy_fix,
+        mask_proxy_credentials,
+    )
+
+    # Apply the fix when the module is imported
+    apply_databricks_proxy_fix()
+except ImportError as e:
+    logger.debug(f"Could not import proxy patch module: {e}")
+
+    # Fallback function for masking credentials
+    def mask_proxy_credentials(url: Optional[str]) -> str:
+        return "***MASKED***" if url else "None"
+
 
 @dataclasses.dataclass
 class TableInfoWithGeneration(TableInfo):
@@ -85,6 +124,32 @@ class QueryFilterWithStatementTypes(QueryFilter):
         return v
 
 
+@dataclasses.dataclass
+class TableUpstream:
+    table_name: str
+    source_type: str
+    last_updated: Optional[datetime] = None
+
+
+@dataclasses.dataclass
+class ExternalUpstream:
+    path: str
+    source_type: str
+    last_updated: Optional[datetime] = None
+
+
+@dataclasses.dataclass
+class TableLineageInfo:
+    upstreams: List[TableUpstream] = dataclasses.field(default_factory=list)
+    external_upstreams: List[ExternalUpstream] = dataclasses.field(default_factory=list)
+    upstream_notebooks: List[NotebookReference] = dataclasses.field(
+        default_factory=list
+    )
+    downstream_notebooks: List[NotebookReference] = dataclasses.field(
+        default_factory=list
+    )
+
+
 class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     _workspace_client: WorkspaceClient
     _workspace_url: str
@@ -98,6 +163,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         warehouse_id: Optional[str],
         report: UnityCatalogReport,
         hive_metastore_proxy: Optional[HiveMetastoreProxy] = None,
+        lineage_data_source: LineageDataSource = LineageDataSource.AUTO,
+        usage_data_source: UsageDataSource = UsageDataSource.AUTO,
+        databricks_api_page_size: int = 0,
     ):
         self._workspace_client = WorkspaceClient(
             host=workspace_url,
@@ -108,9 +176,24 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         self.warehouse_id = warehouse_id or ""
         self.report = report
         self.hive_metastore_proxy = hive_metastore_proxy
+        self.lineage_data_source = lineage_data_source
+        self.usage_data_source = usage_data_source
+        self.databricks_api_page_size = databricks_api_page_size
+        self._sql_connection_params = {
+            "server_hostname": self._workspace_client.config.host.replace(
+                "https://", ""
+            ),
+            "http_path": f"/sql/1.0/warehouses/{self.warehouse_id}",
+            "access_token": self._workspace_client.config.token,
+            "user_agent_entry": "datahub",
+        }
 
     def check_basic_connectivity(self) -> bool:
-        return bool(self._workspace_client.catalogs.list(include_browse=True))
+        return bool(
+            self._workspace_client.catalogs.list(
+                include_browse=True, max_results=self.databricks_api_page_size
+            )
+        )
 
     def assigned_metastore(self) -> Optional[Metastore]:
         response = self._workspace_client.metastores.summary()
@@ -120,7 +203,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         if self.hive_metastore_proxy:
             yield self.hive_metastore_proxy.hive_metastore_catalog(metastore)
 
-        response = self._workspace_client.catalogs.list(include_browse=True)
+        response = self._workspace_client.catalogs.list(
+            include_browse=True, max_results=self.databricks_api_page_size
+        )
         if not response:
             logger.info("Catalogs not found")
             return
@@ -152,7 +237,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             yield from self.hive_metastore_proxy.hive_metastore_schemas(catalog)
             return
         response = self._workspace_client.schemas.list(
-            catalog_name=catalog.name, include_browse=True
+            catalog_name=catalog.name,
+            include_browse=True,
+            max_results=self.databricks_api_page_size,
         )
         if not response:
             logger.info(f"Schemas not found for catalog {catalog.id}")
@@ -174,6 +261,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 catalog_name=schema.catalog.name,
                 schema_name=schema.name,
                 include_browse=True,
+                max_results=self.databricks_api_page_size,
             )
             if not response:
                 logger.info(f"Tables not found for schema {schema.id}")
@@ -189,6 +277,40 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     logger.warning(f"Error parsing table: {e}")
                     self.report.report_warning("table-parse", str(e))
 
+    def ml_models(
+        self, schema: Schema, max_results: Optional[int] = None
+    ) -> Iterable[Model]:
+        response = self._workspace_client.registered_models.list(
+            catalog_name=schema.catalog.name,
+            schema_name=schema.name,
+            max_results=max_results,
+        )
+        for ml_model in response:
+            optional_ml_model = self._create_ml_model(schema, ml_model)
+            if optional_ml_model:
+                yield optional_ml_model
+
+    def ml_model_versions(
+        self, ml_model: Model, include_aliases: bool = False
+    ) -> Iterable[ModelVersion]:
+        response = self._workspace_client.model_versions.list(
+            full_name=ml_model.id,
+            include_browse=True,
+            max_results=self.databricks_api_page_size,
+        )
+        for version in response:
+            if version.version is not None:
+                if include_aliases:
+                    # to get aliases info, use GET
+                    version = self._workspace_client.model_versions.get(
+                        ml_model.id, version.version, include_aliases=True
+                    )
+                optional_ml_model_version = self._create_ml_model_version(
+                    ml_model, version
+                )
+                if optional_ml_model_version:
+                    yield optional_ml_model_version
+
     def service_principals(self) -> Iterable[ServicePrincipal]:
         for principal in self._workspace_client.service_principals.list():
             optional_sp = self._create_service_principal(principal)
@@ -206,7 +328,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         return group_list
 
     def workspace_notebooks(self) -> Iterable[Notebook]:
-        for obj in self._workspace_client.workspace.list("/", recursive=True):
+        workspace_objects_iter = self._workspace_client.workspace.list(
+            "/", recursive=True, max_results=self.databricks_api_page_size
+        )
+        for obj in workspace_objects_iter:
             if obj.object_type == ObjectType.NOTEBOOK and obj.object_id and obj.path:
                 yield Notebook(
                     id=obj.object_id,
@@ -248,7 +373,6 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     def _query_history(
         self,
         filter_by: QueryFilterWithStatementTypes,
-        max_results: int = 1000,
         include_metrics: bool = False,
     ) -> Iterable[QueryInfo]:
         """Manual implementation of the query_history.list() endpoint.
@@ -260,9 +384,10 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
         """
         method = "GET"
         path = "/api/2.0/sql/history/queries"
+
         body: Dict[str, Any] = {
             "include_metrics": include_metrics,
-            "max_results": max_results,  # Max batch size
+            "max_results": self.databricks_api_page_size,  # Max batch size
         }
 
         response: dict = self._workspace_client.api_client.do(  # type: ignore
@@ -280,10 +405,257 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 method, path, body={**body, "page_token": response["next_page_token"]}
             )
 
-    def list_lineages_by_table(
+    def get_query_history_via_system_tables(
+        self,
+        start_time: datetime,
+        end_time: datetime,
+    ) -> Iterable[Query]:
+        """Get query history using system.query.history table.
+
+        This method provides an alternative to the REST API for fetching query history,
+        offering better performance and richer data for large query volumes.
+        """
+        logger.info(
+            f"Fetching query history from system.query.history for period: {start_time} to {end_time}"
+        )
+
+        allowed_types = [typ.value for typ in ALLOWED_STATEMENT_TYPES]
+        statement_type_filter = ", ".join(f"'{typ}'" for typ in allowed_types)
+
+        query = f"""
+            SELECT
+                statement_id,
+                statement_text,
+                statement_type,
+                start_time,
+                end_time,
+                executed_by,
+                executed_as,
+                executed_by_user_id,
+                executed_as_user_id
+            FROM system.query.history
+            WHERE
+                start_time >= %s
+                AND end_time <= %s
+                AND execution_status = 'FINISHED'
+                AND statement_type IN ({statement_type_filter})
+            ORDER BY start_time
+        """
+
+        try:
+            rows = self._execute_sql_query(query, (start_time, end_time))
+            for row in rows:
+                try:
+                    yield Query(
+                        query_id=row.statement_id,
+                        query_text=row.statement_text,
+                        statement_type=(
+                            QueryStatementType(row.statement_type)
+                            if row.statement_type
+                            else None
+                        ),
+                        start_time=row.start_time,
+                        end_time=row.end_time,
+                        user_id=row.executed_by_user_id,
+                        user_name=row.executed_by,
+                        executed_as_user_id=row.executed_as_user_id,
+                        executed_as_user_name=row.executed_as,
+                    )
+                except Exception as e:
+                    logger.warning(f"Error parsing query from system table: {e}")
+                    self.report.report_warning("query-parse-system-table", str(e))
+        except Exception as e:
+            logger.error(
+                f"Error fetching query history from system tables: {e}", exc_info=True
+            )
+            self.report.report_failure(
+                title="Failed to fetch query history from system tables",
+                message="Error querying system.query.history table",
+                context=f"Query period: {start_time} to {end_time}",
+            )
+
+    def _build_datetime_where_conditions(
+        self, start_time: Optional[datetime] = None, end_time: Optional[datetime] = None
+    ) -> str:
+        """Build datetime filtering conditions for lineage queries."""
+        conditions = []
+        if start_time:
+            conditions.append(f"event_time >= '{start_time.isoformat()}'")
+        if end_time:
+            conditions.append(f"event_time <= '{end_time.isoformat()}'")
+        return " AND " + " AND ".join(conditions) if conditions else ""
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_catalog_table_lineage_via_system_tables(
+        self,
+        catalog: str,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ) -> FileBackedDict[TableLineageInfo]:
+        """Get table lineage for all tables in a catalog using system tables."""
+        logger.info(f"Fetching table lineage for catalog: {catalog}")
+        try:
+            additional_where = self._build_datetime_where_conditions(
+                start_time, end_time
+            )
+
+            query = f"""
+                SELECT
+                    entity_type, entity_id,
+                    source_table_full_name, source_type, source_path,
+                    target_table_full_name, target_type,
+                    max(event_time) as last_updated
+                FROM system.access.table_lineage
+                WHERE
+                    (target_table_catalog = %s or source_table_catalog = %s)
+                    {additional_where}
+                GROUP BY
+                    entity_type, entity_id,
+                    source_table_full_name, source_type, source_path,
+                    target_table_full_name, target_type
+                """
+            rows = self._execute_sql_query(query, [catalog, catalog])
+
+            result_dict: FileBackedDict[TableLineageInfo] = FileBackedDict()
+            for row in rows:
+                entity_type = row["entity_type"]
+                entity_id = row["entity_id"]
+                source_full_name = row["source_table_full_name"]
+                target_full_name = row["target_table_full_name"]
+                source_type = row["source_type"]
+                source_path = row["source_path"]
+                last_updated = row["last_updated"]
+
+                # Initialize TableLineageInfo for both source and target tables if they're in our catalog
+                for table_name in [source_full_name, target_full_name]:
+                    if (
+                        table_name
+                        and table_name.startswith(f"{catalog}.")
+                        and table_name not in result_dict
+                    ):
+                        result_dict[table_name] = TableLineageInfo()
+
+                # Process upstream relationships (target table gets upstreams)
+                if target_full_name and target_full_name.startswith(f"{catalog}."):
+                    # Handle table upstreams
+                    if (
+                        source_type in ["TABLE", "VIEW"]
+                        and source_full_name != target_full_name
+                    ):
+                        upstream = TableUpstream(
+                            table_name=source_full_name,
+                            source_type=source_type,
+                            last_updated=last_updated,
+                        )
+                        result_dict[target_full_name].upstreams.append(upstream)
+
+                    # Handle external upstreams (PATH type)
+                    elif source_type == "PATH":
+                        external_upstream = ExternalUpstream(
+                            path=source_path,
+                            source_type=source_type,
+                            last_updated=last_updated,
+                        )
+                        result_dict[target_full_name].external_upstreams.append(
+                            external_upstream
+                        )
+
+                    # Handle upstream notebooks (notebook -> table)
+                    elif entity_type == "NOTEBOOK":
+                        notebook_ref = NotebookReference(
+                            id=entity_id,
+                            last_updated=last_updated,
+                        )
+                        result_dict[target_full_name].upstream_notebooks.append(
+                            notebook_ref
+                        )
+
+                # Process downstream relationships (source table gets downstream notebooks)
+                if (
+                    entity_type == "NOTEBOOK"
+                    and source_full_name
+                    and source_full_name.startswith(f"{catalog}.")
+                ):
+                    notebook_ref = NotebookReference(
+                        id=entity_id,
+                        last_updated=last_updated,
+                    )
+                    result_dict[source_full_name].downstream_notebooks.append(
+                        notebook_ref
+                    )
+
+            return result_dict
+        except Exception as e:
+            logger.warning(
+                f"Error getting table lineage for catalog {catalog}: {e}",
+                exc_info=True,
+            )
+            return FileBackedDict()
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_catalog_column_lineage_via_system_tables(
+        self,
+        catalog: str,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ) -> FileBackedDict[Dict[str, dict]]:
+        """Get column lineage for all tables in a catalog using system tables."""
+        logger.info(f"Fetching column lineage for catalog: {catalog}")
+        try:
+            additional_where = self._build_datetime_where_conditions(
+                start_time, end_time
+            )
+
+            query = f"""
+                SELECT
+                    source_table_catalog, source_table_schema, source_table_name, source_column_name, source_type,
+                    target_table_schema, target_table_name, target_column_name,
+                    max(event_time) as last_updated
+                FROM system.access.column_lineage
+                WHERE
+                    target_table_catalog = %s
+                    AND target_table_schema IS NOT NULL
+                    AND target_table_name IS NOT NULL
+                    AND target_column_name IS NOT NULL
+                    AND source_table_catalog IS NOT NULL
+                    AND source_table_schema IS NOT NULL
+                    AND source_table_name IS NOT NULL
+                    AND source_column_name IS NOT NULL
+                    {additional_where}
+                GROUP BY
+                    source_table_catalog, source_table_schema, source_table_name, source_column_name, source_type,
+                    target_table_schema, target_table_name, target_column_name
+                """
+            rows = self._execute_sql_query(query, [catalog])
+
+            result_dict: FileBackedDict[Dict[str, dict]] = FileBackedDict()
+            for row in rows:
+                result_dict.setdefault(row["target_table_schema"], {}).setdefault(
+                    row["target_table_name"], {}
+                ).setdefault(row["target_column_name"], []).append(
+                    # make fields look like the response from the older HTTP API
+                    {
+                        "catalog_name": row["source_table_catalog"],
+                        "schema_name": row["source_table_schema"],
+                        "table_name": row["source_table_name"],
+                        "name": row["source_column_name"],
+                        "last_updated": row["last_updated"],
+                    }
+                )
+
+            return result_dict
+        except Exception as e:
+            logger.warning(
+                f"Error getting column lineage for catalog {catalog}: {e}",
+                exc_info=True,
+            )
+            return FileBackedDict()
+
+    def list_lineages_by_table_via_http_api(
         self, table_name: str, include_entity_lineage: bool
     ) -> dict:
         """List table lineage by table name."""
+        logger.debug(f"Getting table lineage for {table_name}")
         return self._workspace_client.api_client.do(  # type: ignore
             method="GET",
             path="/api/2.0/lineage-tracking/table-lineage",
@@ -293,67 +665,226 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             },
         )
 
-    def list_lineages_by_column(self, table_name: str, column_name: str) -> dict:
+    def list_lineages_by_column_via_http_api(
+        self, table_name: str, column_name: str
+    ) -> list:
         """List column lineage by table name and column name."""
-        return self._workspace_client.api_client.do(  # type: ignore
-            "GET",
-            "/api/2.0/lineage-tracking/column-lineage",
-            body={"table_name": table_name, "column_name": column_name},
-        )
+        logger.debug(f"Getting column lineage for {table_name}.{column_name}")
+        try:
+            return (
+                self._workspace_client.api_client.do(  # type: ignore
+                    "GET",
+                    "/api/2.0/lineage-tracking/column-lineage",
+                    body={"table_name": table_name, "column_name": column_name},
+                ).get("upstream_cols")
+                or []
+            )
+        except Exception as e:
+            logger.warning(
+                f"Error getting column lineage on table {table_name}, column {column_name}: {e}",
+                exc_info=True,
+            )
+            return []
 
-    def table_lineage(self, table: Table, include_entity_lineage: bool) -> None:
+    def table_lineage(
+        self,
+        table: Table,
+        include_entity_lineage: bool,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ) -> None:
         if table.schema.catalog.type == CustomCatalogType.HIVE_METASTORE_CATALOG:
             # Lineage is not available for Hive Metastore Tables.
             return None
-        # Lineage endpoint doesn't exists on 2.1 version
-        try:
-            response: dict = self.list_lineages_by_table(
-                table_name=table.ref.qualified_table_name,
-                include_entity_lineage=include_entity_lineage,
-            )
 
-            for item in response.get("upstreams") or []:
-                if "tableInfo" in item:
-                    table_ref = TableReference.create_from_lineage(
-                        item["tableInfo"], table.schema.catalog.metastore
-                    )
-                    if table_ref:
-                        table.upstreams[table_ref] = {}
-                elif "fileInfo" in item:
-                    external_ref = ExternalTableReference.create_from_lineage(
-                        item["fileInfo"]
-                    )
-                    if external_ref:
-                        table.external_upstreams.add(external_ref)
-
-                for notebook in item.get("notebookInfos") or []:
-                    table.upstream_notebooks.add(notebook["notebook_id"])
+        try:
+            # Determine lineage data source based on config
+            use_system_tables = False
+            if self.lineage_data_source == LineageDataSource.SYSTEM_TABLES:
+                use_system_tables = True
+            elif self.lineage_data_source == LineageDataSource.API:
+                use_system_tables = False
+            elif self.lineage_data_source == LineageDataSource.AUTO:
+                # Use the newer system tables if we have a SQL warehouse, otherwise fall back
+                # to the older (and slower) HTTP API.
+                use_system_tables = bool(self.warehouse_id)
+            else:
+                assert_never(self.lineage_data_source)
 
-            for item in response.get("downstreams") or []:
-                for notebook in item.get("notebookInfos") or []:
-                    table.downstream_notebooks.add(notebook["notebook_id"])
+            if use_system_tables:
+                self._process_system_table_lineage(table, start_time, end_time)
+            else:
+                self._process_table_lineage_via_http_api(table, include_entity_lineage)
         except Exception as e:
             logger.warning(
                 f"Error getting lineage on table {table.ref}: {e}", exc_info=True
             )
 
-    def get_column_lineage(self, table: Table, column_name: str) -> None:
-        try:
-            response: dict = self.list_lineages_by_column(
-                table_name=table.ref.qualified_table_name,
-                column_name=column_name,
+    def _process_system_table_lineage(
+        self,
+        table: Table,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ) -> None:
+        """Process table lineage using system.access.table_lineage table."""
+        catalog_lineage = self.get_catalog_table_lineage_via_system_tables(
+            table.ref.catalog, start_time, end_time
+        )
+        table_full_name = table.ref.qualified_table_name
+
+        lineage_info = catalog_lineage.get(table_full_name, TableLineageInfo())
+
+        # Process table upstreams
+        for upstream in lineage_info.upstreams:
+            upstream_table_name = upstream.table_name
+            # Parse catalog.schema.table format
+            parts = upstream_table_name.split(".")
+            if len(parts) == 3:
+                catalog_name, schema_name, table_name = parts[0], parts[1], parts[2]
+                table_ref = TableReference(
+                    metastore=table.schema.catalog.metastore.id
+                    if table.schema.catalog.metastore
+                    else None,
+                    catalog=catalog_name,
+                    schema=schema_name,
+                    table=table_name,
+                    last_updated=upstream.last_updated,
+                )
+                table.upstreams[table_ref] = {}
+            else:
+                logger.warning(
+                    f"Unexpected upstream table format: {upstream_table_name} for table {table_full_name}"
+                )
+                continue
+
+        # Process external upstreams
+        for external_upstream in lineage_info.external_upstreams:
+            external_ref = ExternalTableReference(
+                path=external_upstream.path,
+                has_permission=True,
+                name=None,
+                type=None,
+                storage_location=external_upstream.path,
+                last_updated=external_upstream.last_updated,
             )
-            for item in response.get("upstream_cols") or []:
+            table.external_upstreams.add(external_ref)
+
+        # Process upstream notebook lineage
+        for notebook_ref in lineage_info.upstream_notebooks:
+            existing_ref = table.upstream_notebooks.get(notebook_ref.id)
+            if existing_ref is None or (
+                notebook_ref.last_updated
+                and existing_ref.last_updated
+                and notebook_ref.last_updated > existing_ref.last_updated
+            ):
+                table.upstream_notebooks[notebook_ref.id] = notebook_ref
+
+        # Process downstream notebook lineage
+        for notebook_ref in lineage_info.downstream_notebooks:
+            existing_ref = table.downstream_notebooks.get(notebook_ref.id)
+            if existing_ref is None or (
+                notebook_ref.last_updated
+                and existing_ref.last_updated
+                and notebook_ref.last_updated > existing_ref.last_updated
+            ):
+                table.downstream_notebooks[notebook_ref.id] = notebook_ref
+
+    def _process_table_lineage_via_http_api(
+        self, table: Table, include_entity_lineage: bool
+    ) -> None:
+        """Process table lineage using the HTTP API (legacy fallback)."""
+        response: dict = self.list_lineages_by_table_via_http_api(
+            table_name=table.ref.qualified_table_name,
+            include_entity_lineage=include_entity_lineage,
+        )
+
+        for item in response.get("upstreams") or []:
+            if "tableInfo" in item:
                 table_ref = TableReference.create_from_lineage(
-                    item, table.schema.catalog.metastore
+                    item["tableInfo"], table.schema.catalog.metastore
                 )
                 if table_ref:
-                    table.upstreams.setdefault(table_ref, {}).setdefault(
-                        column_name, []
-                    ).append(item["name"])
+                    table.upstreams[table_ref] = {}
+            elif "fileInfo" in item:
+                external_ref = ExternalTableReference.create_from_lineage(
+                    item["fileInfo"]
+                )
+                if external_ref:
+                    table.external_upstreams.add(external_ref)
+
+            for notebook in item.get("notebookInfos") or []:
+                notebook_ref = NotebookReference(
+                    id=notebook["notebook_id"],
+                )
+                table.upstream_notebooks[notebook_ref.id] = notebook_ref
+
+        for item in response.get("downstreams") or []:
+            for notebook in item.get("notebookInfos") or []:
+                notebook_ref = NotebookReference(
+                    id=notebook["notebook_id"],
+                )
+                table.downstream_notebooks[notebook_ref.id] = notebook_ref
+
+    def get_column_lineage(
+        self,
+        table: Table,
+        column_names: List[str],
+        *,
+        max_workers: Optional[int] = None,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ) -> None:
+        try:
+            # Determine lineage data source based on config
+            use_system_tables = False
+            if self.lineage_data_source == LineageDataSource.SYSTEM_TABLES:
+                use_system_tables = True
+            elif self.lineage_data_source == LineageDataSource.API:
+                use_system_tables = False
+            elif self.lineage_data_source == LineageDataSource.AUTO:
+                # Use the newer system tables if we have a SQL warehouse, otherwise fall back
+                # to the older (and slower) HTTP API.
+                use_system_tables = bool(self.warehouse_id)
+            else:
+                assert_never(self.lineage_data_source)
+
+            if use_system_tables:
+                lineage = (
+                    self.get_catalog_column_lineage_via_system_tables(
+                        table.ref.catalog, start_time, end_time
+                    )
+                    .get(table.ref.schema, {})
+                    .get(table.ref.table, {})
+                )
+            else:
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    futures = [
+                        executor.submit(
+                            self.list_lineages_by_column_via_http_api,
+                            table.ref.qualified_table_name,
+                            column_name,
+                        )
+                        for column_name in column_names
+                    ]
+                lineage = {
+                    column_name: future.result()
+                    for column_name, future in zip(column_names, futures)
+                }
+
+            for column_name in column_names:
+                for item in lineage.get(column_name) or []:
+                    table_ref = TableReference.create_from_lineage(
+                        item,
+                        table.schema.catalog.metastore,
+                    )
+                    if table_ref:
+                        table.upstreams.setdefault(table_ref, {}).setdefault(
+                            column_name, []
+                        ).append(item["name"])
+
         except Exception as e:
             logger.warning(
-                f"Error getting column lineage on table {table.ref}, column {column_name}: {e}",
+                f"Error getting column lineage on table {table.ref}: {e}",
                 exc_info=True,
             )
 
@@ -461,6 +992,45 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             if optional_column:
                 yield optional_column
 
+    def _create_ml_model(
+        self, schema: Schema, obj: RegisteredModelInfo
+    ) -> Optional[Model]:
+        if not obj.name or not obj.full_name:
+            self.report.num_ml_models_missing_name += 1
+            return None
+        return Model(
+            id=obj.full_name,
+            name=obj.name,
+            description=obj.comment,
+            schema_name=schema.name,
+            catalog_name=schema.catalog.name,
+            created_at=parse_ts_millis(obj.created_at),
+            updated_at=parse_ts_millis(obj.updated_at),
+        )
+
+    def _create_ml_model_version(
+        self, model: Model, obj: ModelVersionInfo
+    ) -> Optional[ModelVersion]:
+        if obj.version is None:
+            return None
+
+        aliases = []
+        if obj.aliases:
+            for alias in obj.aliases:
+                if alias.alias_name:
+                    aliases.append(alias.alias_name)
+        return ModelVersion(
+            id=f"{model.id}_{obj.version}",
+            name=f"{model.name}_{obj.version}",
+            model=model,
+            version=str(obj.version),
+            aliases=aliases,
+            description=obj.comment,
+            created_at=parse_ts_millis(obj.created_at),
+            updated_at=parse_ts_millis(obj.updated_at),
+            created_by=obj.created_by,
+        )
+
     def _create_service_principal(
         self, obj: DatabricksServicePrincipal
     ) -> Optional[ServicePrincipal]:
@@ -492,3 +1062,176 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             executed_as_user_id=info.executed_as_user_id,
             executed_as_user_name=info.executed_as_user_name,
         )
+
+    def _execute_sql_query(self, query: str, params: Sequence[Any] = ()) -> List[Row]:
+        """Execute SQL query using databricks-sql connector for better performance"""
+        logger.debug(f"Executing SQL query with {len(params)} parameters")
+        if logger.isEnabledFor(logging.DEBUG):
+            # Only log full query in debug mode to avoid performance overhead
+            logger.debug(f"Full SQL query: {query}")
+            if params:
+                logger.debug(f"Query parameters: {params}")
+
+        # Check if warehouse_id is available for SQL operations
+        if not self.warehouse_id:
+            self.report.report_warning(
+                "Cannot execute SQL query",
+                "warehouse_id is not configured. SQL operations require a valid warehouse_id to be set in the Unity Catalog configuration",
+            )
+            logger.warning(
+                "Cannot execute SQL query: warehouse_id is not configured. "
+                "SQL operations require a valid warehouse_id to be set in the Unity Catalog configuration."
+            )
+            return []
+
+        # Log connection parameters (with masked token)
+        masked_params = {**self._sql_connection_params}
+        if "access_token" in masked_params:
+            masked_params["access_token"] = "***MASKED***"
+        logger.debug(f"Using connection parameters: {masked_params}")
+
+        # Log proxy environment variables that affect SQL connections
+        proxy_env_debug = {}
+        for var in ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"]:
+            value = os.environ.get(var)
+            if value:
+                proxy_env_debug[var] = mask_proxy_credentials(value)
+
+        if proxy_env_debug:
+            logger.debug(
+                f"SQL connection will use proxy environment variables: {proxy_env_debug}"
+            )
+        else:
+            logger.debug("No proxy environment variables detected for SQL connection")
+
+        try:
+            with (
+                connect(**self._sql_connection_params) as connection,
+                connection.cursor() as cursor,
+            ):
+                cursor.execute(query, list(params))
+                rows = cursor.fetchall()
+                logger.debug(
+                    f"SQL query executed successfully, returned {len(rows)} rows"
+                )
+                return rows
+
+        except Exception as e:
+            logger.warning(f"Failed to execute SQL query: {e}", exc_info=True)
+            if logger.isEnabledFor(logging.DEBUG):
+                # Only log failed query details in debug mode for security
+                logger.debug(f"SQL query that failed: {query}")
+                logger.debug(f"SQL query parameters: {params}")
+
+            # Check if this might be a proxy-related error
+            error_str = str(e).lower()
+            if any(
+                proxy_keyword in error_str
+                for proxy_keyword in [
+                    "proxy",
+                    "407",
+                    "authentication required",
+                    "tunnel",
+                    "connect",
+                ]
+            ):
+                logger.error(
+                    "SQL query failure appears to be proxy-related. "
+                    "Please check proxy configuration and authentication. "
+                    f"Proxy environment variables detected: {list(proxy_env_debug.keys())}"
+                )
+
+            return []
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_schema_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching schema tags for catalog: `{catalog}`")
+
+        query = f"SELECT * FROM `{catalog}`.information_schema.schema_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, tag_name, tag_value = row
+            schema_key = f"{catalog_name}.{schema_name}"
+
+            if schema_key not in result_dict:
+                result_dict[schema_key] = []
+
+            result_dict[schema_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_catalog_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
+
+        query = f"SELECT * FROM `{catalog}`.information_schema.catalog_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, tag_name, tag_value = row
+
+            if catalog_name not in result_dict:
+                result_dict[catalog_name] = []
+
+            result_dict[catalog_name].append(
+                UnityCatalogTag(key=tag_name, value=tag_value)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_table_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
+
+        query = f"SELECT * FROM `{catalog}`.information_schema.table_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, table_name, tag_name, tag_value = row
+            table_key = f"{catalog_name}.{schema_name}.{table_name}"
+
+            if table_key not in result_dict:
+                result_dict[table_key] = []
+
+            result_dict[table_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value if tag_value else None)
+            )
+
+        return result_dict
+
+    @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))
+    def get_column_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
+        """Optimized version using databricks-sql"""
+        logger.info(f"Fetching column tags for catalog: `{catalog}`")
+
+        query = f"SELECT * FROM `{catalog}`.information_schema.column_tags"
+        rows = self._execute_sql_query(query)
+
+        result_dict: Dict[str, List[UnityCatalogTag]] = {}
+
+        for row in rows:
+            catalog_name, schema_name, table_name, column_name, tag_name, tag_value = (
+                row
+            )
+            column_key = f"{catalog_name}.{schema_name}.{table_name}.{column_name}"
+
+            if column_key not in result_dict:
+                result_dict[column_key] = []
+
+            result_dict[column_key].append(
+                UnityCatalogTag(key=tag_name, value=tag_value if tag_value else None)
+            )
+
+        return result_dict