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

datahub/ingestion/source/unity/config.py

@@ -2,23 +2,25 @@ import logging
 import os
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List, Optional, Union
-from urllib.parse import urlparse
 
 import pydantic
 from pydantic import Field
 from typing_extensions import Literal
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import (
+    AllowDenyPattern,
+    ConfigEnum,
+    ConfigModel,
+    HiddenFromDocs,
+)
 from datahub.configuration.source_common import (
     DatasetSourceConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
 )
 from datahub.configuration.validate_field_removal import pydantic_removed_field
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
-from datahub.ingestion.source.ge_data_profiler import DATABRICKS
 from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig
-from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StatefulStaleMetadataRemovalConfig,
 )
@@ -26,6 +28,7 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
     StatefulProfilingConfigMixin,
 )
+from datahub.ingestion.source.unity.connection import UnityCatalogConnectionConfig
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 from datahub.ingestion.source_config.operation_config import (
     OperationConfig,
@@ -35,6 +38,22 @@ from datahub.utilities.global_warning_util import add_global_warning
 
 logger = logging.getLogger(__name__)
 
+# Configuration default constants
+INCLUDE_TAGS_DEFAULT = True
+INCLUDE_HIVE_METASTORE_DEFAULT = True
+
+
+class LineageDataSource(ConfigEnum):
+    AUTO = "AUTO"
+    SYSTEM_TABLES = "SYSTEM_TABLES"
+    API = "API"
+
+
+class UsageDataSource(ConfigEnum):
+    AUTO = "AUTO"
+    SYSTEM_TABLES = "SYSTEM_TABLES"
+    API = "API"
+
 
 class UnityCatalogProfilerConfig(ConfigModel):
     method: str = Field(
@@ -118,6 +137,7 @@ class UnityCatalogGEProfilerConfig(UnityCatalogProfilerConfig, GEProfilingConfig
 
 
 class UnityCatalogSourceConfig(
+    UnityCatalogConnectionConfig,
     SQLCommonConfig,
     StatefulIngestionConfigBase,
     BaseUsageConfig,
@@ -125,23 +145,6 @@ class UnityCatalogSourceConfig(
     StatefulProfilingConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
 ):
-    token: str = pydantic.Field(description="Databricks personal access token")
-    workspace_url: str = pydantic.Field(
-        description="Databricks workspace url. e.g. https://my-workspace.cloud.databricks.com"
-    )
-    warehouse_id: Optional[str] = pydantic.Field(
-        default=None,
-        description="SQL Warehouse id, for running queries. If not set, will use the default warehouse.",
-    )
-    include_hive_metastore: bool = pydantic.Field(
-        default=True,
-        description="Whether to ingest legacy `hive_metastore` catalog. This requires executing queries on SQL warehouse.",
-    )
-    workspace_name: Optional[str] = pydantic.Field(
-        default=None,
-        description="Name of the workspace. Default to deployment name present in workspace_url",
-    )
-
     include_metastore: bool = pydantic.Field(
         default=False,
         description=(
@@ -229,6 +232,15 @@ class UnityCatalogSourceConfig(
         description="Option to enable/disable ownership generation for metastores, catalogs, schemas, and tables.",
     )
 
+    include_tags: bool = pydantic.Field(
+        default=INCLUDE_TAGS_DEFAULT,
+        description=(
+            "Option to enable/disable column/table tag extraction. "
+            "Requires warehouse_id to be set since tag extraction needs to query system.information_schema.tags. "
+            "If warehouse_id is not provided, this will be automatically disabled to allow ingestion to continue."
+        ),
+    )
+
     _rename_table_ownership = pydantic_renamed_field(
         "include_table_ownership", "include_ownership"
     )
@@ -238,15 +250,40 @@ class UnityCatalogSourceConfig(
         description="Option to enable/disable lineage generation. Currently we have to call a rest call per column to get column level lineage due to the Databrick api which can slow down ingestion. ",
     )
 
+    lineage_data_source: LineageDataSource = pydantic.Field(
+        default=LineageDataSource.AUTO,
+        description=(
+            "Source for lineage data extraction. Options: "
+            f"'{LineageDataSource.AUTO.value}' - Use system tables when SQL warehouse is available, fallback to API; "
+            f"'{LineageDataSource.SYSTEM_TABLES.value}' - Force use of system.access.table_lineage and system.access.column_lineage tables (requires SQL warehouse); "
+            f"'{LineageDataSource.API.value}' - Force use of REST API endpoints for lineage data"
+        ),
+    )
+
+    ignore_start_time_lineage: bool = pydantic.Field(
+        default=False,
+        description="Option to ignore the start_time and retrieve all available lineage. When enabled, the start_time filter will be set to zero to extract all lineage events regardless of the configured time window.",
+    )
+
     column_lineage_column_limit: int = pydantic.Field(
         default=300,
         description="Limit the number of columns to get column level lineage. ",
     )
 
-    lineage_max_workers: int = pydantic.Field(
+    lineage_max_workers: HiddenFromDocs[int] = pydantic.Field(
         default=5 * (os.cpu_count() or 4),
         description="Number of worker threads to use for column lineage thread pool executor. Set to 1 to disable.",
-        hidden_from_docs=True,
+    )
+
+    databricks_api_page_size: int = pydantic.Field(
+        default=0,
+        ge=0,
+        description=(
+            "Page size for Databricks API calls when listing resources (catalogs, schemas, tables, etc.). "
+            "When set to 0 (default), uses server-side configured page length (recommended). "
+            "When set to a positive value, the page length is the minimum of this value and the server configured value. "
+            "Must be a non-negative integer."
+        ),
     )
 
     include_usage_statistics: bool = Field(
@@ -254,6 +291,17 @@ class UnityCatalogSourceConfig(
         description="Generate usage statistics.",
     )
 
+    usage_data_source: UsageDataSource = pydantic.Field(
+        default=UsageDataSource.AUTO,
+        description=(
+            "Source for usage/query history data extraction. Options: "
+            f"'{UsageDataSource.AUTO.value}' (default) - Automatically use system.query.history table when SQL warehouse is configured, otherwise fall back to REST API. "
+            "This provides better performance for multi-workspace setups and large query volumes when warehouse_id is set. "
+            f"'{UsageDataSource.SYSTEM_TABLES.value}' - Force use of system.query.history table (requires SQL warehouse and SELECT permission on system.query.history). "
+            f"'{UsageDataSource.API.value}' - Force use of REST API endpoints for query history (legacy method, may have limitations with multiple workspaces)."
+        ),
+    )
+
     # TODO: Remove `type:ignore` by refactoring config
     profiling: Union[
         UnityCatalogGEProfilerConfig, UnityCatalogAnalyzeProfilerConfig
@@ -273,19 +321,68 @@ class UnityCatalogSourceConfig(
         description="Details about the delta lake, incase to emit siblings",
     )
 
-    scheme: str = DATABRICKS
+    include_ml_model_aliases: bool = pydantic.Field(
+        default=False,
+        description="Whether to include ML model aliases in the ingestion.",
+    )
+
+    ml_model_max_results: int = pydantic.Field(
+        default=1000,
+        ge=0,
+        description="Maximum number of ML models to ingest.",
+    )
+
+    _forced_disable_tag_extraction: bool = pydantic.PrivateAttr(default=False)
+    _forced_disable_hive_metastore_extraction = pydantic.PrivateAttr(default=False)
+
+    include_hive_metastore: bool = pydantic.Field(
+        default=INCLUDE_HIVE_METASTORE_DEFAULT,
+        description="Whether to ingest legacy `hive_metastore` catalog. This requires executing queries on SQL warehouse.",
+    )
 
-    def get_sql_alchemy_url(self, database: Optional[str] = None) -> str:
-        uri_opts = {"http_path": f"/sql/1.0/warehouses/{self.warehouse_id}"}
-        if database:
-            uri_opts["catalog"] = database
-        return make_sqlalchemy_uri(
-            scheme=self.scheme,
-            username="token",
-            password=self.token,
-            at=urlparse(self.workspace_url).netloc,
-            db=database,
-            uri_opts=uri_opts,
+    workspace_name: Optional[str] = pydantic.Field(
+        default=None,
+        description="Name of the workspace. Default to deployment name present in workspace_url",
+    )
+
+    def __init__(self, **data):
+        # First, let the parent handle the root validators and field processing
+        super().__init__(**data)
+
+        # After model creation, check if we need to auto-disable features
+        # based on the final warehouse_id value (which may have been set by root validators)
+        include_tags_original = data.get("include_tags", INCLUDE_TAGS_DEFAULT)
+        include_hive_metastore_original = data.get(
+            "include_hive_metastore", INCLUDE_HIVE_METASTORE_DEFAULT
+        )
+
+        # Track what we're force-disabling
+        forced_disable_tag_extraction = False
+        forced_disable_hive_metastore_extraction = False
+
+        # Check if features should be auto-disabled based on final warehouse_id
+        if include_tags_original and not self.warehouse_id:
+            forced_disable_tag_extraction = True
+            self.include_tags = False  # Modify the model attribute directly
+            logger.warning(
+                "warehouse_id is not set but include_tags=True. "
+                "Automatically disabling tag extraction since it requires SQL queries. "
+                "Set warehouse_id to enable tag extraction."
+            )
+
+        if include_hive_metastore_original and not self.warehouse_id:
+            forced_disable_hive_metastore_extraction = True
+            self.include_hive_metastore = False  # Modify the model attribute directly
+            logger.warning(
+                "warehouse_id is not set but include_hive_metastore=True. "
+                "Automatically disabling hive metastore extraction since it requires SQL queries. "
+                "Set warehouse_id to enable hive metastore extraction."
+            )
+
+        # Set private attributes
+        self._forced_disable_tag_extraction = forced_disable_tag_extraction
+        self._forced_disable_hive_metastore_extraction = (
+            forced_disable_hive_metastore_extraction
         )
 
     def is_profiling_enabled(self) -> bool:
@@ -344,11 +441,6 @@ class UnityCatalogSourceConfig(
                 "When `warehouse_id` is set, it must match the `warehouse_id` in `profiling`."
             )
 
-        if values.get("include_hive_metastore") and not values.get("warehouse_id"):
-            raise ValueError(
-                "When `include_hive_metastore` is set, `warehouse_id` must be set."
-            )
-
         if values.get("warehouse_id") and profiling and not profiling.warehouse_id:
             profiling.warehouse_id = values["warehouse_id"]
 
@@ -357,6 +449,34 @@ class UnityCatalogSourceConfig(
 
         return values
 
+    @pydantic.root_validator(skip_on_failure=True)
+    def validate_lineage_data_source_with_warehouse(
+        cls, values: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        lineage_data_source = values.get("lineage_data_source", LineageDataSource.AUTO)
+        warehouse_id = values.get("warehouse_id")
+
+        if lineage_data_source == LineageDataSource.SYSTEM_TABLES and not warehouse_id:
+            raise ValueError(
+                f"lineage_data_source='{LineageDataSource.SYSTEM_TABLES.value}' requires warehouse_id to be set"
+            )
+
+        return values
+
+    @pydantic.root_validator(skip_on_failure=True)
+    def validate_usage_data_source_with_warehouse(
+        cls, values: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        usage_data_source = values.get("usage_data_source", UsageDataSource.AUTO)
+        warehouse_id = values.get("warehouse_id")
+
+        if usage_data_source == UsageDataSource.SYSTEM_TABLES and not warehouse_id:
+            raise ValueError(
+                f"usage_data_source='{UsageDataSource.SYSTEM_TABLES.value}' requires warehouse_id to be set"
+            )
+
+        return values
+
     @pydantic.validator("schema_pattern", always=True)
     def schema_pattern_should__always_deny_information_schema(
         cls, v: AllowDenyPattern

datahub/ingestion/source/unity/connection.py

@@ -0,0 +1,61 @@
+"""Databricks Unity Catalog connection configuration."""
+
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+
+import pydantic
+from pydantic import Field
+
+from datahub.configuration.common import ConfigModel
+from datahub.ingestion.source.sql.sqlalchemy_uri import make_sqlalchemy_uri
+
+DATABRICKS = "databricks"
+
+
+class UnityCatalogConnectionConfig(ConfigModel):
+    """
+    Configuration for connecting to Databricks Unity Catalog.
+    Contains only connection-related fields that can be reused across different sources.
+    """
+
+    scheme: str = DATABRICKS
+    token: str = pydantic.Field(description="Databricks personal access token")
+    workspace_url: str = pydantic.Field(
+        description="Databricks workspace url. e.g. https://my-workspace.cloud.databricks.com"
+    )
+    warehouse_id: Optional[str] = pydantic.Field(
+        default=None,
+        description=(
+            "SQL Warehouse id, for running queries. Must be explicitly provided to enable SQL-based features. "
+            "Required for the following features that need SQL access: "
+            "1) Tag extraction (include_tags=True) - queries system.information_schema.tags "
+            "2) Hive Metastore catalog (include_hive_metastore=True) - queries legacy hive_metastore catalog "
+            "3) System table lineage (lineage_data_source=SYSTEM_TABLES) - queries system.access.table_lineage/column_lineage "
+            "4) Data profiling (profiling.enabled=True) - runs SELECT/ANALYZE queries on tables. "
+            "When warehouse_id is missing, these features will be automatically disabled (with warnings) to allow ingestion to continue."
+        ),
+    )
+
+    extra_client_options: Dict[str, Any] = Field(
+        default={},
+        description="Additional options to pass to Databricks SQLAlchemy client.",
+    )
+
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+
+    def get_sql_alchemy_url(self, database: Optional[str] = None) -> str:
+        uri_opts = {"http_path": f"/sql/1.0/warehouses/{self.warehouse_id}"}
+        if database:
+            uri_opts["catalog"] = database
+        return make_sqlalchemy_uri(
+            scheme=self.scheme,
+            username="token",
+            password=self.token,
+            at=urlparse(self.workspace_url).netloc,
+            db=database,
+            uri_opts=uri_opts,
+        )
+
+    def get_options(self) -> dict:
+        return self.extra_client_options

datahub/ingestion/source/unity/connection_test.py

@@ -19,6 +19,7 @@ class UnityCatalogConnectionTest:
             self.config.token,
             self.config.profiling.warehouse_id,
             report=self.report,
+            databricks_api_page_size=self.config.databricks_api_page_size,
         )
 
     def get_connection_test(self) -> TestConnectionReport:

datahub/ingestion/source/unity/platform_resource_repository.py

@@ -0,0 +1,19 @@
+import logging
+
+from datahub.api.entities.external.external_entities import (
+    PlatformResourceRepository,
+)
+from datahub.ingestion.source.unity.tag_entities import (
+    UnityCatalogTagPlatformResource,
+    UnityCatalogTagPlatformResourceId,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class UnityCatalogPlatformResourceRepository(
+    PlatformResourceRepository[
+        UnityCatalogTagPlatformResourceId, UnityCatalogTagPlatformResource
+    ]
+):
+    """Unity Catalog-specific platform resource repository with tag-related operations."""