acryl-datahub 1.2.0.7rc3__py3-none-any.whl → 1.2.0.8__py3-none-any.whl

datahub/ingestion/source/unity/config.py

@@ -35,6 +35,10 @@ 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"
@@ -137,10 +141,18 @@ class UnityCatalogSourceConfig(
     )
     warehouse_id: Optional[str] = pydantic.Field(
         default=None,
-        description="SQL Warehouse id, for running queries. If not set, will use the default warehouse.",
+        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."
+        ),
     )
     include_hive_metastore: bool = pydantic.Field(
-        default=True,
+        default=INCLUDE_HIVE_METASTORE_DEFAULT,
         description="Whether to ingest legacy `hive_metastore` catalog. This requires executing queries on SQL warehouse.",
     )
     workspace_name: Optional[str] = pydantic.Field(
@@ -236,8 +248,12 @@ class UnityCatalogSourceConfig(
     )
 
     include_tags: bool = pydantic.Field(
-        default=True,
-        description="Option to enable/disable column/table tag extraction.",
+        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(
@@ -310,8 +326,6 @@ 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.",
@@ -323,6 +337,51 @@ class UnityCatalogSourceConfig(
         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)
+
+    scheme: str = DATABRICKS
+
+    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 get_sql_alchemy_url(self, database: Optional[str] = None) -> str:
         uri_opts = {"http_path": f"/sql/1.0/warehouses/{self.warehouse_id}"}
         if database:
@@ -392,11 +451,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"]
 

datahub/ingestion/source/unity/proxy.py

@@ -4,6 +4,7 @@ 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, Sequence, Union, cast
@@ -71,6 +72,23 @@ logger: logging.Logger = logging.getLogger(__name__)
 _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):
     generation: Optional[int] = None
@@ -411,7 +429,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             query = f"""
                 SELECT
                     entity_type, entity_id,
-                    source_table_full_name, source_type,
+                    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
@@ -420,7 +438,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     {additional_where}
                 GROUP BY
                     entity_type, entity_id,
-                    source_table_full_name, source_type,
+                    source_table_full_name, source_type, source_path,
                     target_table_full_name, target_type
                 """
             rows = self._execute_sql_query(query, [catalog, catalog])
@@ -432,6 +450,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 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
@@ -460,7 +479,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     # Handle external upstreams (PATH type)
                     elif source_type == "PATH":
                         external_upstream = ExternalUpstream(
-                            path=source_full_name,
+                            path=source_path,
                             source_type=source_type,
                             last_updated=last_updated,
                         )
@@ -973,16 +992,82 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
 
     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))
-                return cursor.fetchall()
+                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}")
+            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))

datahub/ingestion/source/unity/proxy_patch.py

@@ -0,0 +1,321 @@
+"""
+Proxy authentication patch for databricks-sql < 3.0 compatibility.
+
+This module provides proxy authentication fixes for databricks-sql connector < 3.0
+to resolve "407 Proxy Authentication Required" errors that occur even when
+proxy environment variables are correctly set.
+
+The patch implements the same fix as Databricks PR #354:
+https://github.com/databricks/databricks-sql-python/pull/354
+"""
+
+import logging
+import os
+import urllib.parse
+from typing import Dict, Optional
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+PROXY_VARS = ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"]
+
+
+def mask_proxy_credentials(url: Optional[str]) -> str:
+    """Mask credentials in proxy URL for safe logging."""
+    if not url:
+        return "None"
+
+    try:
+        parsed = urllib.parse.urlparse(url)
+        if parsed.username:
+            # Replace credentials with masked version
+            masked_netloc = parsed.netloc
+            if parsed.username and parsed.password:
+                masked_netloc = masked_netloc.replace(
+                    f"{parsed.username}:{parsed.password}@", f"{parsed.username}:***@"
+                )
+            elif parsed.username:
+                masked_netloc = masked_netloc.replace(
+                    f"{parsed.username}@", f"{parsed.username}:***@"
+                )
+
+            return urllib.parse.urlunparse(
+                (
+                    parsed.scheme,
+                    masked_netloc,
+                    parsed.path,
+                    parsed.params,
+                    parsed.query,
+                    parsed.fragment,
+                )
+            )
+        else:
+            return url
+    except Exception:
+        return "***INVALID_URL***"
+
+
+def _ensure_thrift_imports():
+    """Ensure required thrift imports are loaded before accessing thrift_http_client.
+
+    The databricks-sql thrift_http_client requires thrift.transport.THttpClient.THttpClient
+    to be accessible. This is achieved by importing the required modules in the right order.
+    """
+    try:
+        # Import thrift submodules - this makes them accessible as attributes
+        import thrift.transport.THttpClient  # noqa: F401 # Used to make thrift.transport accessible
+
+        logger.debug("Successfully imported required thrift modules")
+    except Exception as e:
+        logger.debug(f"Could not import thrift modules: {e}")
+        raise
+
+
+def _log_proxy_environment():
+    """Log detected proxy environment variables for debugging."""
+    proxy_env_vars = {}
+    for var in PROXY_VARS:
+        value = os.environ.get(var)
+        if value:
+            masked_value = mask_proxy_credentials(value)
+            proxy_env_vars[var] = masked_value
+
+    if proxy_env_vars:
+        logger.info(f"Detected proxy environment variables: {proxy_env_vars}")
+    else:
+        logger.debug("No proxy environment variables detected")
+
+
+def _basic_proxy_auth_header(proxy_url: str) -> Optional[Dict[str, str]]:
+    """Create proxy authentication header using the same method as Databricks >= 3.0.
+
+    Based on the basic_proxy_auth_header method from databricks-sql-connector >= 3.0:
+    https://github.com/databricks/databricks-sql-python/pull/354
+    """
+    try:
+        from urllib3.util import make_headers
+
+        parsed = urllib.parse.urlparse(proxy_url)
+        if parsed.username and parsed.password:
+            # Code reused from https://github.com/databricks/databricks-sql-python/pull/354
+            # URL decode the username and password (same as Databricks method)
+            username = urllib.parse.unquote(parsed.username)
+            password = urllib.parse.unquote(parsed.password)
+            auth_string = f"{username}:{password}"
+
+            # Create proxy URL without credentials
+            proxy_host_port = f"{parsed.scheme}://{parsed.hostname}"
+            if parsed.port:
+                proxy_host_port += f":{parsed.port}"
+
+            # Code reused from https://github.com/databricks/databricks-sql-python/pull/354
+            # Use make_headers like the newer Databricks version does
+            proxy_headers = make_headers(proxy_basic_auth=auth_string)
+
+            return {
+                "proxy_url": proxy_host_port,
+                "proxy_headers": proxy_headers,
+                "auth_string": auth_string,  # Keep for backward compatibility with tests
+            }
+    except Exception as e:
+        logger.debug(f"Failed to create proxy auth header from URL {proxy_url}: {e}")
+
+    return None
+
+
+def _handle_proxy_connection(self, original_open, pool_kwargs):
+    """Handle proxy connection setup with authentication headers."""
+    from urllib3.poolmanager import ProxyManager
+
+    logger.info(f"Using proxy for connection to {self.host}:{self.port}")
+    proxy_uri = getattr(self, "proxy_uri", None)
+    logger.debug(
+        f"Proxy URI: {mask_proxy_credentials(proxy_uri) if proxy_uri else 'None'}"
+    )
+
+    # Compute proxy authentication headers properly (the bug fix!)
+    proxy_headers = None
+    proxy_env_found = None
+    for env_var in ["HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy"]:
+        proxy_url = os.environ.get(env_var)
+        if proxy_url:
+            logger.debug(
+                f"Found proxy URL in {env_var}: {mask_proxy_credentials(proxy_url)}"
+            )
+            auth_info = _basic_proxy_auth_header(proxy_url)
+            if auth_info:
+                proxy_headers = auth_info["proxy_headers"]
+                proxy_env_found = env_var
+                logger.debug(f"Successfully created proxy headers from {env_var}")
+                break
+            else:
+                logger.debug(
+                    f"No authentication info found in proxy URL from {env_var}"
+                )
+
+    if proxy_headers:
+        logger.info(f"Using proxy authentication headers from {proxy_env_found}")
+    else:
+        logger.warning(
+            "No proxy authentication headers could be created from environment variables"
+        )
+
+    proxy_manager = ProxyManager(
+        self.proxy_uri,
+        num_pools=1,
+        proxy_headers=proxy_headers,
+    )
+
+    # Validate proxy manager attributes
+    if not hasattr(self, "realhost") or not hasattr(self, "realport"):
+        logger.warning(
+            "THttpClient missing realhost/realport attributes, falling back to original"
+        )
+        return original_open(self)
+
+    # Set up the connection pool
+    self._THttpClient__pool = proxy_manager.connection_from_host(
+        host=self.realhost,
+        port=self.realport,
+        scheme=self.scheme,
+        pool_kwargs=pool_kwargs,  # type: ignore
+    )
+    logger.debug(f"Created proxy connection pool for {self.realhost}:{self.realport}")
+
+
+def _create_patched_open_method(original_open):
+    """Create the patched THttpClient.open method with proxy authentication fix."""
+
+    def patched_open(self):
+        """Patched version of THttpClient.open following databricks-sql >= 3.0 structure.
+
+        This is largely copied from the >= 3.0 implementation:
+        https://github.com/databricks/databricks-sql-python/pull/354/files
+        """
+        logger.debug(
+            f"Patched THttpClient.open called for host={getattr(self, 'host', 'unknown')}, scheme={getattr(self, 'scheme', 'unknown')}"
+        )
+
+        try:
+            # Validate required attributes
+            required_attrs = ["scheme", "host", "port", "max_connections"]
+            missing_attrs = [attr for attr in required_attrs if not hasattr(self, attr)]
+            if missing_attrs:
+                logger.warning(
+                    f"THttpClient missing required attributes: {missing_attrs}, falling back to original"
+                )
+                return original_open(self)
+
+            # Code structure reused from https://github.com/databricks/databricks-sql-python/pull/354
+            # Determine pool class based on scheme
+            if self.scheme == "http":
+                from urllib3 import HTTPConnectionPool
+
+                pool_class = HTTPConnectionPool
+            elif self.scheme == "https":
+                from urllib3 import HTTPSConnectionPool
+
+                pool_class = HTTPSConnectionPool
+            else:
+                logger.warning(
+                    f"Unknown scheme '{self.scheme}', falling back to original"
+                )
+                return original_open(self)
+
+            _pool_kwargs = {"maxsize": self.max_connections}
+            logger.debug(f"Pool kwargs: {_pool_kwargs}")
+
+            if self.using_proxy():
+                return _handle_proxy_connection(self, original_open, _pool_kwargs)
+            else:
+                logger.debug(f"Direct connection (no proxy) to {self.host}:{self.port}")
+                self._THttpClient__pool = pool_class(
+                    self.host, self.port, **_pool_kwargs
+                )
+
+            logger.debug("Patched THttpClient.open completed successfully")
+
+        except Exception as e:
+            logger.warning(
+                f"Error in proxy auth patch: {e}, falling back to original",
+                exc_info=True,
+            )
+            # Fallback to original implementation
+            try:
+                return original_open(self)
+            except Exception as fallback_error:
+                logger.error(
+                    f"Fallback to original THttpClient.open also failed: {fallback_error}",
+                    exc_info=True,
+                )
+                raise
+
+    return patched_open
+
+
+def apply_databricks_proxy_fix():
+    """Apply the databricks-sql < 3.0 proxy authentication fix at module import time.
+
+    This implements the same fix as Databricks PR #354 to resolve
+    "407 Proxy Authentication Required" errors that occur even when
+    all proxy environment variables are correctly set.
+
+    Note: This fix may not work with all thrift versions due to compatibility issues
+    between databricks-sql-connector 2.9.6 and newer thrift versions. The fix will
+    gracefully fail with a warning if thrift compatibility issues are detected.
+    The main SQL functionality will continue to work normally without this fix.
+    """
+    _log_proxy_environment()
+    logger.info("Applying databricks-sql proxy authentication fix...")
+
+    try:
+        _ensure_thrift_imports()
+        import databricks.sql.auth.thrift_http_client as thrift_http
+
+        # Store original method for fallback
+        original_open = getattr(thrift_http.THttpClient, "open", None)
+        if not original_open:
+            logger.warning("Could not find THttpClient.open method to patch")
+            return False
+
+        logger.debug(f"Found THttpClient.open method at {original_open}")
+
+        # Apply the patch
+        patched_open = _create_patched_open_method(original_open)
+        thrift_http.THttpClient.open = patched_open
+        logger.info("Successfully applied databricks-sql proxy authentication fix")
+
+        # Verify the patch was applied
+        current_method = getattr(thrift_http.THttpClient, "open", None)
+        if current_method == patched_open:
+            logger.debug(
+                "Patch verification successful: THttpClient.open is now the patched version"
+            )
+            return True
+        else:
+            logger.warning(
+                "Patch verification failed: THttpClient.open was not replaced correctly"
+            )
+            return False
+
+    except ImportError as e:
+        logger.debug(f"Could not import databricks-sql internals for proxy patch: {e}")
+        return False
+    except AttributeError as e:
+        if "thrift" in str(e).lower() and "transport" in str(e).lower():
+            warning_msg = (
+                f"Databricks-sql proxy authentication patch could not be applied due to thrift version incompatibility: {e}. "
+                "In most environments, the SQL connection will still work without this patch."
+            )
+            logger.warning(warning_msg)
+            # Import here to avoid circular imports
+            from datahub.utilities.global_warning_util import add_global_warning
+
+            add_global_warning(warning_msg)
+        else:
+            logger.error(
+                f"Failed to apply databricks-sql proxy patch: {e}", exc_info=True
+            )
+        return False
+    except Exception as e:
+        logger.error(f"Failed to apply databricks-sql proxy patch: {e}", exc_info=True)
+        return False

datahub/ingestion/source/unity/source.py

@@ -257,6 +257,18 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         else:
             self.platform_resource_repository = None
 
+        if self.config._forced_disable_tag_extraction:
+            self.report.report_warning(
+                "Some features disabled because of configuration conflicts",
+                "Tag Extraction is disabled due to missing warehouse_id in config",
+            )
+
+        if self.config._forced_disable_hive_metastore_extraction:
+            self.report.report_warning(
+                "Some features disabled because of configuration conflicts",
+                "Hive Metastore Extraction is disabled due to missing warehouse_id in config",
+            )
+
         # Include platform resource repository in report for automatic cache statistics
         if self.config.include_tags and self.platform_resource_repository:
             self.report.tag_urn_resolver_cache = self.platform_resource_repository

datahub/ingestion/source/usage/usage_common.py

@@ -268,6 +268,7 @@ class UsageAggregator(Generic[ResourceType]):
             user,
             query,
             fields,
+            user_email_pattern=self.config.user_email_pattern,
             count=count,
         )