openmetadata-ingestion 1.9.0.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/ingestion/source/database/databricks/auth.py

@@ -0,0 +1,99 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This module provides authentication utilities for Databricks and Unity Catalog connections.
+"""
+from typing import Union
+
+from databricks.sdk.core import Config, azure_service_principal, oauth_service_principal
+
+from metadata.generated.schema.entity.services.connections.database.databricks.azureAdSetup import (
+    AzureAdSetup,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.databricksOAuth import (
+    DatabricksOauth,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
+    PersonalAccessToken,
+)
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
+
+
+def get_personal_access_token_auth(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+) -> dict:
+    """
+    Configure Personal Access Token authentication
+    """
+    return {"access_token": connection.authType.token.get_secret_value()}
+
+
+def get_databricks_oauth_auth(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+):
+    """
+    Create Databricks OAuth2 M2M credentials provider for Service Principal authentication
+    """
+
+    def credential_provider():
+        hostname = connection.hostPort.split(":")[0]
+        config = Config(
+            host=f"https://{hostname}",
+            client_id=connection.authType.clientId,
+            client_secret=connection.authType.clientSecret.get_secret_value(),
+        )
+        return oauth_service_principal(config)
+
+    return {"credentials_provider": credential_provider}
+
+
+def get_azure_ad_auth(connection: Union[DatabricksConnection, UnityCatalogConnection]):
+    """
+    Create Azure AD credentials provider for Azure Service Principal authentication
+    """
+
+    def credential_provider():
+        hostname = connection.hostPort.split(":")[0]
+        config = Config(
+            host=f"https://{hostname}",
+            azure_client_secret=connection.authType.azureClientSecret.get_secret_value(),
+            azure_client_id=connection.authType.azureClientId,
+            azure_tenant_id=connection.authType.azureTenantId,
+        )
+        return azure_service_principal(config)
+
+    return {"credentials_provider": credential_provider}
+
+
+def get_auth_config(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+) -> dict:
+    """
+    Get authentication configuration for Databricks connection
+    """
+    auth_method = {
+        PersonalAccessToken: get_personal_access_token_auth,
+        DatabricksOauth: get_databricks_oauth_auth,
+        AzureAdSetup: get_azure_ad_auth,
+    }.get(type(connection.authType))
+
+    if not auth_method:
+        raise ValueError(
+            f"Unsupported authentication type: {type(connection.authType)}"
+        )
+
+    return auth_method(connection)

metadata/ingestion/source/database/databricks/client.py

@@ -11,8 +11,10 @@
 """
 Client to interact with databricks apis
 """
+import base64
 import json
 import traceback
+from collections import defaultdict
 from datetime import timedelta
 from typing import Iterable, List, Optional, Tuple, Union
 
@@ -61,24 +63,31 @@ class DatabricksClient:
     ):
         self.config = config
         base_url, *_ = self.config.hostPort.split(":")
-        auth_token = self.config.token.get_secret_value()
         self.base_url = f"https://{base_url}{API_VERSION}"
         self.base_query_url = f"{self.base_url}{QUERIES_PATH}"
         self.base_job_url = f"https://{base_url}{JOB_API_VERSION}/jobs"
         self.jobs_list_url = f"{self.base_job_url}/list"
         self.jobs_run_list_url = f"{self.base_job_url}/runs/list"
         self.headers = {
-            "Authorization": f"Bearer {auth_token}",
+            **self._get_auth_header(),
             "Content-Type": "application/json",
         }
         self.api_timeout = self.config.connectionTimeout or 120
-        self.job_table_lineage: dict[str, list[dict[str, str]]] = {}
+        self._job_table_lineage_executed: bool = False
+        self.job_table_lineage: dict[str, list[dict[str, str]]] = defaultdict(list)
+        self._job_column_lineage_executed: bool = False
         self.job_column_lineage: dict[
             str, dict[Tuple[str, str], list[Tuple[str, str]]]
-        ] = {}
+        ] = defaultdict(lambda: defaultdict(list))
         self.engine = engine
         self.client = requests
 
+    def _get_auth_header(self) -> dict[str, str]:
+        """
+        Method to get auth header
+        """
+        return {"Authorization": f"Bearer {self.config.token.get_secret_value()}"}
+
     def test_query_api_access(self) -> None:
         res = self.client.get(
             self.base_query_url, headers=self.headers, timeout=self.api_timeout
@@ -88,12 +97,23 @@ class DatabricksClient:
 
     def test_lineage_query(self) -> None:
         try:
+            lookback_days = getattr(self.config, "lineageLookBackDays", 90)
             with self.engine.connect() as connection:
                 test_table_lineage = connection.execute(
-                    text(DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB + " LIMIT 1")
+                    text(
+                        DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB.format(
+                            lookback_days=lookback_days
+                        )
+                        + " LIMIT 1"
+                    )
                 )
                 test_column_lineage = connection.execute(
-                    text(DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB + " LIMIT 1")
+                    text(
+                        DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB.format(
+                            lookback_days=lookback_days
+                        )
+                        + " LIMIT 1"
+                    )
                 )
                 # Check if queries executed successfully by fetching results
                 table_result = test_table_lineage.fetchone()
@@ -193,7 +213,6 @@ class DatabricksClient:
         """
         Method returns List all the created jobs in a Databricks Workspace
         """
-        self.cache_lineage()
         try:
             iteration_count = 1
             data = {"limit": PAGE_SIZE, "expand_tasks": True, "offset": 0}
@@ -263,10 +282,19 @@ class DatabricksClient:
 
     def get_table_lineage(self, job_id: str) -> List[dict[str, str]]:
         """
-        Method returns table lineage for a job by the specified job_id
+        Method returns table lineage for a job by the specified job_id.
+        On first call, eagerly fetches ALL job lineage in bulk for optimal performance.
         """
         try:
-            return self.job_table_lineage.get(str(job_id))
+            if not self._job_table_lineage_executed:
+                logger.info(
+                    "First lineage request detected - performing bulk lineage fetch for all jobs"
+                )
+                self.cache_lineage()
+
+            # Return cached lineage for this specific job
+            return self.job_table_lineage.get(str(job_id), [])
+
         except Exception as exc:
             logger.debug(
                 f"Error getting table lineage for job {job_id} due to {traceback.format_exc()}"
@@ -281,7 +309,12 @@ class DatabricksClient:
         Method returns column lineage for a job by the specified job_id and table key
         """
         try:
-            return self.job_column_lineage.get(str(job_id), {}).get(TableKey)
+            if not self._job_column_lineage_executed:
+                logger.debug("Job column lineage not found. Executing cache_lineage...")
+                self.cache_lineage()
+
+            return self.job_column_lineage.get(str(job_id), {}).get(TableKey, [])
+
         except Exception as exc:
             logger.debug(
                 f"Error getting column lineage for table {TableKey} due to {traceback.format_exc()}"
@@ -305,53 +338,170 @@ class DatabricksClient:
 
     def cache_lineage(self):
         """
-        Method caches table and column lineage for a job by the specified job_id
+        Method caches table and column lineage for ALL jobs.
         """
-        logger.info(f"Caching table lineage")
-        table_lineage = self.run_lineage_query(DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB)
-        if table_lineage:
-            for row in table_lineage:
-                try:
-                    if row.job_id not in self.job_table_lineage:
-                        self.job_table_lineage[row.job_id] = []
-                    self.job_table_lineage[row.job_id].append(
-                        {
-                            "source_table_full_name": row.source_table_full_name,
-                            "target_table_full_name": row.target_table_full_name,
-                        }
-                    )
-                except Exception as exc:
-                    logger.debug(
-                        f"Error parsing row: {row} due to {traceback.format_exc()}"
-                    )
-                    continue
+        lookback_days = getattr(self.config, "lineageLookBackDays", 90)
+        logger.info(f"Caching table lineage (lookback: {lookback_days} days)")
+        table_lineage = self.run_lineage_query(
+            DATABRICKS_GET_TABLE_LINEAGE_FOR_JOB.format(lookback_days=lookback_days)
+        )
+        for row in table_lineage or []:
+            try:
+                self.job_table_lineage[row.job_id].append(
+                    {
+                        "source_table_full_name": row.source_table_full_name,
+                        "target_table_full_name": row.target_table_full_name,
+                    }
+                )
+            except Exception as exc:
+                logger.debug(
+                    f"Error parsing row: {row} due to {traceback.format_exc()}"
+                )
+                continue
+        self._job_table_lineage_executed = True
 
         # Not every job has column lineage, so we need to check if the job exists in the column_lineage table
         # we will cache the column lineage for jobs that have column lineage
-        logger.info("Caching column lineage")
-        column_lineage = self.run_lineage_query(DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB)
-        if column_lineage:
-            for row in column_lineage:
-                try:
-                    table_key = (
-                        row.source_table_full_name,
-                        row.target_table_full_name,
-                    )
-                    column_pair = (
-                        row.source_column_name,
-                        row.target_column_name,
-                    )
+        logger.info(f"Caching column lineage (lookback: {lookback_days} days)")
+        column_lineage = self.run_lineage_query(
+            DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB.format(lookback_days=lookback_days)
+        )
+        for row in column_lineage or []:
+            try:
+                table_key = (
+                    row.source_table_full_name,
+                    row.target_table_full_name,
+                )
+                column_pair = (
+                    row.source_column_name,
+                    row.target_column_name,
+                )
+
+                self.job_column_lineage[row.job_id][table_key].append(column_pair)
 
-                    if row.job_id not in self.job_column_lineage:
-                        self.job_column_lineage[row.job_id] = {}
+            except Exception as exc:
+                logger.debug(
+                    f"Error parsing row: {row} due to {traceback.format_exc()}"
+                )
+                continue
+        self._job_column_lineage_executed = True
+        logger.debug("Table and column lineage caching completed.")
 
-                    if table_key not in self.job_column_lineage[row.job_id]:
-                        self.job_column_lineage[row.job_id][table_key] = []
+    def get_pipeline_details(self, pipeline_id: str) -> Optional[dict]:
+        """
+        Get DLT pipeline configuration including libraries and notebooks
+        """
+        try:
+            url = f"{self.base_url}/pipelines/{pipeline_id}"
+            response = self.client.get(
+                url,
+                headers=self.headers,
+                timeout=self.api_timeout,
+            )
+            if response.status_code == 200:
+                return response.json()
+            logger.warning(
+                f"Failed to get pipeline details for {pipeline_id}: {response.status_code}"
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error getting pipeline details for {pipeline_id}: {exc}")
+        return None
 
-                    self.job_column_lineage[row.job_id][table_key].append(column_pair)
+    def list_pipelines(self) -> Iterable[dict]:
+        """
+        List all DLT (Delta Live Tables) pipelines in the workspace
+        Uses the Pipelines API (/api/2.0/pipelines)
+        """
+        try:
+            url = f"{self.base_url}/pipelines"
+            params = {"max_results": PAGE_SIZE}
 
-                except Exception as exc:
-                    logger.debug(
-                        f"Error parsing row: {row} due to {traceback.format_exc()}"
+            response = self.client.get(
+                url,
+                params=params,
+                headers=self.headers,
+                timeout=self.api_timeout,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                pipelines = data.get("statuses", [])
+                logger.info(f"Found {len(pipelines)} DLT pipelines")
+                yield from pipelines
+
+                # Handle pagination if there's a next_page_token
+                while data.get("next_page_token"):
+                    params["page_token"] = data["next_page_token"]
+                    response = self.client.get(
+                        url,
+                        params=params,
+                        headers=self.headers,
+                        timeout=self.api_timeout,
                     )
-                    continue
+                    if response.status_code == 200:
+                        data = response.json()
+                        yield from data.get("statuses", [])
+                    else:
+                        break
+            else:
+                logger.warning(
+                    f"Failed to list pipelines: {response.status_code} - {response.text}"
+                )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error listing DLT pipelines: {exc}")
+
+    def list_workspace_objects(self, path: str) -> List[dict]:
+        """
+        List objects in a Databricks workspace directory
+        """
+        try:
+            url = f"{self.base_url}/workspace/list"
+            params = {"path": path}
+
+            response = self.client.get(
+                url,
+                params=params,
+                headers=self.headers,
+                timeout=self.api_timeout,
+            )
+
+            if response.status_code == 200:
+                return response.json().get("objects", [])
+            else:
+                logger.warning(
+                    f"Failed to list workspace directory {path}: {response.text}"
+                )
+                return []
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error listing workspace directory {path}: {exc}")
+            return []
+
+    def export_notebook_source(self, notebook_path: str) -> Optional[str]:
+        """
+        Export notebook source code from Databricks workspace
+        """
+        try:
+            url = f"{self.base_url}/workspace/export"
+            params = {"path": notebook_path, "format": "SOURCE"}
+
+            response = self.client.get(
+                url,
+                params=params,
+                headers=self.headers,
+                timeout=self.api_timeout,
+            )
+
+            if response.status_code == 200:
+                content = response.json().get("content")
+                if content:
+                    return base64.b64decode(content).decode("utf-8")
+            logger.warning(
+                f"Failed to export notebook {notebook_path}: {response.status_code}"
+            )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Error exporting notebook {notebook_path}: {exc}")
+        return None

metadata/ingestion/source/database/databricks/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from functools import partial
 from typing import Optional
 
@@ -38,9 +39,17 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.databricks.auth import get_auth_config
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
     DATABRICKS_SQL_STATEMENT_TEST,
+    TEST_CATALOG_TAGS,
+    TEST_COLUMN_LINEAGE,
+    TEST_COLUMN_TAGS,
+    TEST_SCHEMA_TAGS,
+    TEST_TABLE_LINEAGE,
+    TEST_TABLE_TAGS,
+    TEST_VIEW_DEFINITIONS,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -56,11 +65,17 @@ class DatabricksEngineWrapper:
         self.inspector = inspect(engine)
         self.schemas = None
         self.first_schema = None
+        self.first_catalog = None
 
-    def get_schemas(self):
+    def get_schemas(self, schema_name: Optional[str] = None):
         """Get schemas and cache them"""
+        if schema_name is not None:
+            with self.engine.connect() as connection:
+                connection.execute(f"USE CATALOG `{self.first_catalog}`")
+            self.first_schema = schema_name
+            return [schema_name]
         if self.schemas is None:
-            self.schemas = self.inspector.get_schema_names()
+            self.schemas = self.inspector.get_schema_names(database=self.first_catalog)
             if self.schemas:
                 # Find the first schema that's not a system schema
                 for schema in self.schemas:
@@ -81,7 +96,11 @@ class DatabricksEngineWrapper:
         if self.first_schema is None:
             self.get_schemas()  # This will set first_schema
         if self.first_schema:
-            return self.inspector.get_table_names(self.first_schema)
+            with self.engine.connect() as connection:
+                tables = connection.execute(
+                    f"SHOW TABLES IN `{self.first_catalog}`.`{self.first_schema}`"
+                )
+            return tables
         return []
 
     def get_views(self):
@@ -89,13 +108,30 @@ class DatabricksEngineWrapper:
         if self.first_schema is None:
             self.get_schemas()  # This will set first_schema
         if self.first_schema:
-            return self.inspector.get_view_names(self.first_schema)
+            with self.engine.connect() as connection:
+                views = connection.execute(
+                    f"SHOW VIEWS IN `{self.first_catalog}`.`{self.first_schema}`"
+                )
+            return views
         return []
 
+    def get_catalogs(self, catalog_name: Optional[str] = None):
+        """Get catalogs"""
+        catalogs = []
+        if catalog_name is not None:
+            self.first_catalog = catalog_name
+            return [catalog_name]
+        with self.engine.connect() as connection:
+            catalogs = connection.execute(DATABRICKS_GET_CATALOGS).fetchall()
+            for catalog in catalogs:
+                if catalog[0] != "__databricks_internal":
+                    self.first_catalog = catalog[0]
+                    break
+        return catalogs
+
 
 def get_connection_url(connection: DatabricksConnection) -> str:
-    url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
-    return url
+    return f"{connection.scheme.value}://{connection.hostPort}"
 
 
 def get_connection(connection: DatabricksConnection) -> Engine:
@@ -103,17 +139,27 @@ def get_connection(connection: DatabricksConnection) -> Engine:
     Create connection
     """
 
+    if not connection.connectionArguments:
+        connection.connectionArguments = init_empty_connection_arguments()
+
     if connection.httpPath:
-        if not connection.connectionArguments:
-            connection.connectionArguments = init_empty_connection_arguments()
         connection.connectionArguments.root["http_path"] = connection.httpPath
 
-    return create_generic_db_connection(
+    auth_args = get_auth_config(connection)
+
+    original_connection_arguments = connection.connectionArguments
+    connection.connectionArguments = deepcopy(original_connection_arguments)
+    connection.connectionArguments.root.update(auth_args)
+
+    engine = create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
 
+    connection.connectionArguments = original_connection_arguments
+    return engine
+
 
 def test_connection(
     metadata: OpenMetadata,
@@ -142,15 +188,20 @@ def test_connection(
     # Create wrapper to avoid multiple schema calls
     engine_wrapper = DatabricksEngineWrapper(connection)
 
+    # Helper function to get first catalog for tag queries
+    def get_first_catalog():
+        catalogs = engine_wrapper.get_catalogs(catalog_name=service_connection.catalog)
+        return catalogs[0] if catalogs else service_connection.catalog or "main"
+
     test_fn = {
         "CheckAccess": partial(test_connection_engine_step, connection),
-        "GetSchemas": engine_wrapper.get_schemas,
+        "GetSchemas": partial(
+            engine_wrapper.get_schemas, schema_name=service_connection.databaseSchema
+        ),
         "GetTables": engine_wrapper.get_tables,
         "GetViews": engine_wrapper.get_views,
         "GetDatabases": partial(
-            test_database_query,
-            engine=connection,
-            statement=DATABRICKS_GET_CATALOGS,
+            engine_wrapper.get_catalogs, catalog_name=service_connection.catalog
         ),
         "GetQueries": partial(
             test_database_query,
@@ -159,6 +210,41 @@ def test_connection(
                 query_history=service_connection.queryHistoryTable
             ),
         ),
+        "GetViewDefinitions": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_VIEW_DEFINITIONS,
+        ),
+        "GetCatalogTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_CATALOG_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetSchemaTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_SCHEMA_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetTableTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_TABLE_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetColumnTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_COLUMN_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetTableLineage": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_TABLE_LINEAGE,
+        ),
+        "GetColumnLineage": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_COLUMN_LINEAGE,
+        ),
     }
 
     return test_connection_steps(

metadata/ingestion/source/database/databricks/data_diff/__init__.py

@@ -0,0 +1 @@
+"""Unity Catalog data diff module"""

metadata/ingestion/source/database/databricks/data_diff/data_diff.py

@@ -0,0 +1,12 @@
+"""Databricks spec for data diff - uses Databricks connection"""
+
+from metadata.ingestion.source.database.common.data_diff.databricks_base import (
+    DatabricksBaseTableParameter,
+)
+
+
+class DatabricksTableParameter(DatabricksBaseTableParameter):
+    """Databricks table parameter setter - uses Databricks connection
+    which is databricks-based for data diff operations"""
+
+    pass