Add kerberos auth support

Signed-off-by: Vikrant Puppala <[email protected]>

Author: vikrantpuppala

poetry.lock

@@ -26,6 +26,7 @@ pyarrow = [
     { version = ">=18.0.0", python = ">=3.13", optional=true }
 ]
 pyjwt = "^2.0.0"
+requests-kerberos = "^0.15.0"
 
 
 [tool.poetry.extras]

pyproject.toml

@@ -46,9 +46,7 @@ def __init__(
         retry_stop_after_attempts_duration: Optional[float] = None,
         retry_delay_default: Optional[float] = None,
         retry_dangerous_codes: Optional[List[int]] = None,
-        http_proxy: Optional[str] = None,
-        proxy_username: Optional[str] = None,
-        proxy_password: Optional[str] = None,
+        proxy_auth_method: Optional[str] = None,
         pool_connections: Optional[int] = None,
         pool_maxsize: Optional[int] = None,
         user_agent: Optional[str] = None,
@@ -79,9 +77,7 @@ def __init__(
         )
         self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
-        self.http_proxy = http_proxy
-        self.proxy_username = proxy_username
-        self.proxy_password = proxy_password
+        self.proxy_auth_method = proxy_auth_method
         self.pool_connections = pool_connections or 10
         self.pool_maxsize = pool_maxsize or 20
         self.user_agent = user_agent

src/databricks/sql/auth/common.py

@@ -32,6 +32,7 @@ def __init__(
         ssl_options: Optional[SSLOptions] = None,
         max_connections: int = 1,
         retry_policy: Union[DatabricksRetryPolicy, int] = 0,
+        **kwargs,
     ):
         self._ssl_options = ssl_options
 
@@ -63,7 +64,10 @@ def __init__(
                 self.path += "?%s" % parsed.query
 
         # Handle proxy settings using shared utility
-        proxy_uri, proxy_auth = detect_and_parse_proxy(self.scheme, self.host)
+        proxy_auth_method = kwargs.get("_proxy_auth_method")
+        proxy_uri, proxy_auth = detect_and_parse_proxy(
+            self.scheme, self.host, proxy_auth_method=proxy_auth_method
+        )
 
         if proxy_uri:
             parsed_proxy = urllib.parse.urlparse(proxy_uri)

src/databricks/sql/auth/thrift_http_client.py

@@ -125,7 +125,10 @@ def __init__(
             self.retry_policy = 0
 
         # Handle proxy settings using shared utility
-        proxy_uri, proxy_auth = detect_and_parse_proxy(self.scheme, self.host)
+        proxy_auth_method = kwargs.get("_proxy_auth_method")
+        proxy_uri, proxy_auth = detect_and_parse_proxy(
+            self.scheme, self.host, proxy_auth_method=proxy_auth_method
+        )
 
         if proxy_uri:
             parsed_proxy = urllib.parse.urlparse(proxy_uri)

src/databricks/sql/backend/sea/utils/http_client.py

@@ -191,6 +191,12 @@ def __init__(
         self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
 
         additional_transport_args = {}
+
+        # Add proxy authentication method if specified
+        proxy_auth_method = kwargs.get("_proxy_auth_method")
+        if proxy_auth_method:
+            additional_transport_args["_proxy_auth_method"] = proxy_auth_method
+
         _max_redirects: Union[None, int] = kwargs.get("_retry_max_redirects")
 
         if _max_redirects:

src/databricks/sql/backend/thrift_backend.py

@@ -1,6 +1,7 @@
 import ssl
 import urllib.parse
 import urllib.request
+import logging
 from typing import Dict, Any, Optional, Tuple, Union
 
 from urllib3 import HTTPConnectionPool, HTTPSConnectionPool, ProxyManager
@@ -9,9 +10,14 @@
 from databricks.sql.auth.retry import DatabricksRetryPolicy
 from databricks.sql.types import SSLOptions
 
+logger = logging.getLogger(__name__)
+
 
 def detect_and_parse_proxy(
-    scheme: str, host: str = None, skip_bypass: bool = False
+    scheme: str,
+    host: str = None,
+    skip_bypass: bool = False,
+    proxy_auth_method: Optional[str] = None,
 ) -> Tuple[Optional[str], Optional[Dict[str, str]]]:
     """
     Detect system proxy and return proxy URI and headers using standardized logic.
@@ -20,6 +26,7 @@ def detect_and_parse_proxy(
         scheme: URL scheme (http/https)
         host: Target hostname (optional, only needed for bypass checking)
         skip_bypass: If True, skip proxy bypass checking and return proxy config if found
+        proxy_auth_method: Authentication method ('basic', 'negotiate', or None)
 
     Returns:
         Tuple of (proxy_uri, proxy_headers) or (None, None) if no proxy
@@ -40,10 +47,41 @@ def detect_and_parse_proxy(
         return None, None
 
     parsed_proxy = urllib.parse.urlparse(proxy)
-    proxy_headers = create_basic_proxy_auth_headers(parsed_proxy)
+
+    # Generate appropriate auth headers based on method
+    if proxy_auth_method == "negotiate":
+        proxy_headers = _generate_negotiate_headers(parsed_proxy.hostname)
+    elif proxy_auth_method == "basic" or proxy_auth_method is None:
+        # Default to basic if method not specified (backward compatibility)
+        proxy_headers = create_basic_proxy_auth_headers(parsed_proxy)
+    else:
+        raise ValueError(f"Unsupported proxy_auth_method: {proxy_auth_method}")
+
     return proxy, proxy_headers
 
 
+def _generate_negotiate_headers(proxy_hostname: str) -> Optional[Dict[str, str]]:
+    """Generate Kerberos/SPNEGO authentication headers"""
+    try:
+        from requests_kerberos import HTTPKerberosAuth
+
+        logger.debug(
+            f"Attempting to generate Kerberos SPNEGO token for proxy: {proxy_hostname}"
+        )
+        auth = HTTPKerberosAuth()
+        negotiate_details = auth.generate_request_header(
+            None, proxy_hostname, is_preemptive=True
+        )
+        if negotiate_details:
+            return {"proxy-authorization": negotiate_details}
+        else:
+            logger.debug("Unable to generate kerberos proxy auth headers")
+    except Exception as e:
+        logger.error(f"Error generating Kerberos proxy auth headers: {e}")
+
+    return None
+
+
 def create_basic_proxy_auth_headers(parsed_proxy) -> Optional[Dict[str, str]]:
     """
     Create basic auth headers for proxy if credentials are provided.

src/databricks/sql/common/http_utils.py

@@ -15,7 +15,6 @@
 from databricks.sql.common.http import HttpMethod
 from databricks.sql.common.http_utils import (
     detect_and_parse_proxy,
-    create_basic_proxy_auth_headers,
 )
 
 logger = logging.getLogger(__name__)
@@ -124,7 +123,9 @@ def _setup_pool_managers(self):
         try:
             # Use shared proxy detection logic, skipping bypass since we handle that per-request
             proxy_url, proxy_auth = detect_and_parse_proxy(
-                self.scheme, skip_bypass=True
+                self.scheme,
+                skip_bypass=True,
+                proxy_auth_method=self.config.proxy_auth_method,
             )
 
             if proxy_url:

src/databricks/sql/common/unified_http_client.py

@@ -919,9 +919,7 @@ def build_client_context(server_hostname: str, version: str, **kwargs):
         ),
         retry_delay_default=kwargs.get("_retry_delay_default"),
         retry_dangerous_codes=kwargs.get("_retry_dangerous_codes"),
-        http_proxy=kwargs.get("_http_proxy"),
-        proxy_username=kwargs.get("_proxy_username"),
-        proxy_password=kwargs.get("_proxy_password"),
+        proxy_auth_method=kwargs.get("_proxy_auth_method"),
         pool_connections=kwargs.get("_pool_connections"),
         pool_maxsize=kwargs.get("_pool_maxsize"),
     )