databricks-sql-connector 4.1.3__tar.gz → 4.1.5__tar.gz

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Release History
 
+# 4.1.5 (2026-03-23)
+- Add _retry_server_directed_only mode for Retry-After header compliance (databricks/databricks-sql-python#757 by @sd-db)
+
+# 4.1.4 (2025-10-15)
+- Add support for Token Federation (databricks/databricks-sql-python#691 by @madhav-db)
+- Add metric view support (databricks/databricks-sql-python#688 by @shivam2680)
+- Increased time limit for long running queries (databricks/databricks-sql-python#686 by @jprakash-db)
+
 # 4.1.3 (2025-09-17)
 - Query tags integration (databricks/databricks-sql-python#663 by @sreekanth-db)
 - Add variant support (databricks/databricks-sql-python#560 by @shivam2680)

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/PKG-INFO

@@ -1,9 +1,8 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.3
 Name: databricks-sql-connector
-Version: 4.1.3
+Version: 4.1.5
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
-License-File: LICENSE
 Author: Databricks
 Author-email: databricks-sql-connector-maintainers@databricks.com
 Requires-Python: >=3.8.0,<4.0.0
@@ -15,7 +14,6 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
-Classifier: Programming Language :: Python :: 3.14
 Provides-Extra: pyarrow
 Requires-Dist: lz4 (>=4.0.2,<5.0.0)
 Requires-Dist: oauthlib (>=3.1.0,<4.0.0)

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.1.3"
+version = "4.1.5"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/__init__.py

@@ -68,7 +68,7 @@ DATETIME = DBAPITypeObject("timestamp")
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.1.3"
+__version__ = "4.1.5"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/auth/auth.py

@@ -8,13 +8,17 @@ from databricks.sql.auth.authenticators import (
     AzureServicePrincipalCredentialProvider,
 )
 from databricks.sql.auth.common import AuthType, ClientContext
+from databricks.sql.auth.token_federation import TokenFederationProvider
 
 
 def get_auth_provider(cfg: ClientContext, http_client):
+    # Determine the base auth provider
+    base_provider: Optional[AuthProvider] = None
+
     if cfg.credentials_provider:
-        return ExternalAuthProvider(cfg.credentials_provider)
+        base_provider = ExternalAuthProvider(cfg.credentials_provider)
     elif cfg.auth_type == AuthType.AZURE_SP_M2M.value:
-        return ExternalAuthProvider(
+        base_provider = ExternalAuthProvider(
             AzureServicePrincipalCredentialProvider(
                 cfg.hostname,
                 cfg.azure_client_id,
@@ -29,7 +33,7 @@ def get_auth_provider(cfg: ClientContext, http_client):
         assert cfg.oauth_client_id is not None
         assert cfg.oauth_scopes is not None
 
-        return DatabricksOAuthProvider(
+        base_provider = DatabricksOAuthProvider(
             cfg.hostname,
             cfg.oauth_persistence,
             cfg.oauth_redirect_port_range,
@@ -39,17 +43,17 @@ def get_auth_provider(cfg: ClientContext, http_client):
             cfg.auth_type,
         )
     elif cfg.access_token is not None:
-        return AccessTokenAuthProvider(cfg.access_token)
+        base_provider = AccessTokenAuthProvider(cfg.access_token)
     elif cfg.use_cert_as_auth and cfg.tls_client_cert_file:
         # no op authenticator. authentication is performed using ssl certificate outside of headers
-        return AuthProvider()
+        base_provider = AuthProvider()
     else:
         if (
             cfg.oauth_redirect_port_range is not None
             and cfg.oauth_client_id is not None
             and cfg.oauth_scopes is not None
         ):
-            return DatabricksOAuthProvider(
+            base_provider = DatabricksOAuthProvider(
                 cfg.hostname,
                 cfg.oauth_persistence,
                 cfg.oauth_redirect_port_range,
@@ -61,6 +65,17 @@ def get_auth_provider(cfg: ClientContext, http_client):
         else:
             raise RuntimeError("No valid authentication settings!")
 
+    # Always wrap with token federation (falls back gracefully if not needed)
+    if base_provider:
+        return TokenFederationProvider(
+            hostname=cfg.hostname,
+            external_provider=base_provider,
+            http_client=http_client,
+            identity_federation_client_id=cfg.identity_federation_client_id,
+        )
+
+    return base_provider
+
 
 PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]
 PYSQL_OAUTH_CLIENT_ID = "databricks-sql-python"
@@ -114,5 +129,6 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs)
         else redirect_port_range,
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
+        identity_federation_client_id=kwargs.get("identity_federation_client_id"),
     )
     return get_auth_provider(cfg, http_client)

databricks_sql_connector-4.1.5/src/databricks/sql/auth/auth_utils.py

@@ -0,0 +1,64 @@
+import logging
+import jwt
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Tuple
+from urllib.parse import urlparse
+
+logger = logging.getLogger(__name__)
+
+
+def parse_hostname(hostname: str) -> str:
+    """
+    Normalize the hostname to include scheme and trailing slash.
+
+    Args:
+        hostname: The hostname to normalize
+
+    Returns:
+        Normalized hostname with scheme and trailing slash
+    """
+    if not hostname.startswith("http://") and not hostname.startswith("https://"):
+        hostname = f"https://{hostname}"
+    if not hostname.endswith("/"):
+        hostname = f"{hostname}/"
+    return hostname
+
+
+def decode_token(access_token: str) -> Optional[Dict]:
+    """
+    Decode a JWT token without verification to extract claims.
+
+    Args:
+        access_token: The JWT access token to decode
+
+    Returns:
+        Decoded token claims or None if decoding fails
+    """
+    try:
+        return jwt.decode(access_token, options={"verify_signature": False})
+    except Exception as e:
+        logger.debug("Failed to decode JWT token: %s", e)
+        return None
+
+
+def is_same_host(url1: str, url2: str) -> bool:
+    """
+    Check if two URLs have the same host.
+
+    Args:
+        url1: First URL
+        url2: Second URL
+
+    Returns:
+        True if hosts are the same, False otherwise
+    """
+    try:
+        host1 = urlparse(url1).netloc
+        host2 = urlparse(url2).netloc
+        # Handle port differences (e.g., example.com vs example.com:443)
+        host1_without_port = host1.split(":")[0]
+        host2_without_port = host2.split(":")[0]
+        return host1_without_port == host2_without_port
+    except Exception as e:
+        logger.debug("Failed to parse URLs: %s", e)
+        return False

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/auth/common.py

@@ -37,6 +37,7 @@ class ClientContext:
         tls_client_cert_file: Optional[str] = None,
         oauth_persistence=None,
         credentials_provider=None,
+        identity_federation_client_id: Optional[str] = None,
         # HTTP client configuration parameters
         ssl_options=None,  # SSLOptions type
         socket_timeout: Optional[float] = None,
@@ -46,6 +47,7 @@ class ClientContext:
         retry_stop_after_attempts_duration: Optional[float] = None,
         retry_delay_default: Optional[float] = None,
         retry_dangerous_codes: Optional[List[int]] = None,
+        respect_server_retry_after_header: Optional[bool] = None,
         proxy_auth_method: Optional[str] = None,
         pool_connections: Optional[int] = None,
         pool_maxsize: Optional[int] = None,
@@ -65,6 +67,7 @@ class ClientContext:
         self.tls_client_cert_file = tls_client_cert_file
         self.oauth_persistence = oauth_persistence
         self.credentials_provider = credentials_provider
+        self.identity_federation_client_id = identity_federation_client_id
 
         # HTTP client configuration
         self.ssl_options = ssl_options
@@ -77,6 +80,7 @@ class ClientContext:
         )
         self.retry_delay_default = retry_delay_default or 5.0
         self.retry_dangerous_codes = retry_dangerous_codes or []
+        self.respect_server_retry_after_header = bool(respect_server_retry_after_header)
         self.proxy_auth_method = proxy_auth_method
         self.pool_connections = pool_connections or 10
         self.pool_maxsize = pool_maxsize or 20

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/auth/retry.py

@@ -94,6 +94,7 @@ class DatabricksRetryPolicy(Retry):
         stop_after_attempts_duration: float,
         delay_default: float,
         force_dangerous_codes: List[int],
+        respect_server_retry_after_header: bool = False,
         urllib3_kwargs: dict = {},
     ):
         # These values do not change from one command to the next
@@ -103,6 +104,7 @@ class DatabricksRetryPolicy(Retry):
         self.stop_after_attempts_duration = stop_after_attempts_duration
         self._delay_default = delay_default
         self.force_dangerous_codes = force_dangerous_codes
+        self.respect_server_retry_after_header = respect_server_retry_after_header
 
         # the urllib3 kwargs are a mix of configuration (some of which we override)
         # and counters like `total` or `connect` which may change between successive retries
@@ -202,6 +204,7 @@ class DatabricksRetryPolicy(Retry):
             stop_after_attempts_duration=self.stop_after_attempts_duration,
             delay_default=self.delay_default,
             force_dangerous_codes=self.force_dangerous_codes,
+            respect_server_retry_after_header=self.respect_server_retry_after_header,
             urllib3_kwargs={},
         )
 
@@ -323,7 +326,9 @@ class DatabricksRetryPolicy(Retry):
 
         return proposed_backoff
 
-    def should_retry(self, method: str, status_code: int) -> Tuple[bool, str]:
+    def should_retry(
+        self, method: str, status_code: int, has_retry_after: bool = False
+    ) -> Tuple[bool, str]:
         """This method encapsulates the connector's approach to retries.
 
         We always retry a request unless one of these conditions is met:
@@ -381,6 +386,15 @@ class DatabricksRetryPolicy(Retry):
         if not self._is_method_retryable(method):
             return False, "Only POST requests are retried"
 
+        # When respect_server_retry_after_header is enabled, only retry when the
+        # server explicitly signals it's safe via a Retry-After header. This prevents
+        # duplicate side effects for non-idempotent operations.
+        if self.respect_server_retry_after_header and not has_retry_after:
+            return (
+                False,
+                "respect_server_retry_after_header mode: no Retry-After header present",
+            )
+
         # Request failed with 404 and was a GetOperationStatus. This is not recoverable. Don't retry.
         if status_code == 404 and self.command_type == CommandType.GET_OPERATION_STATUS:
             return (
@@ -450,7 +464,7 @@ class DatabricksRetryPolicy(Retry):
         Logs a debug message if the request will be retried
         """
 
-        should_retry, msg = self.should_retry(method, status_code)
+        should_retry, msg = self.should_retry(method, status_code, has_retry_after)
 
         if should_retry:
             logger.debug(msg)

databricks_sql_connector-4.1.5/src/databricks/sql/auth/token_federation.py

@@ -0,0 +1,206 @@
+import logging
+import json
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Tuple
+from urllib.parse import urlencode
+
+from databricks.sql.auth.authenticators import AuthProvider
+from databricks.sql.auth.auth_utils import (
+    parse_hostname,
+    decode_token,
+    is_same_host,
+)
+from databricks.sql.common.http import HttpMethod
+
+logger = logging.getLogger(__name__)
+
+
+class Token:
+    """
+    Represents an OAuth token with expiration management.
+    """
+
+    def __init__(self, access_token: str, token_type: str = "Bearer"):
+        """
+        Initialize a token.
+
+        Args:
+            access_token: The access token string
+            token_type: The token type (default: Bearer)
+        """
+        self.access_token = access_token
+        self.token_type = token_type
+        self.expiry_time = self._calculate_expiry()
+
+    def _calculate_expiry(self) -> datetime:
+        """
+        Calculate the token expiry time from JWT claims.
+
+        Returns:
+            The token expiry datetime
+        """
+        decoded = decode_token(self.access_token)
+        if decoded and "exp" in decoded:
+            # Use JWT exp claim with 1 minute buffer
+            return datetime.fromtimestamp(decoded["exp"]) - timedelta(minutes=1)
+        # Default to 1 hour if no expiry info
+        return datetime.now() + timedelta(hours=1)
+
+    def is_expired(self) -> bool:
+        """
+        Check if the token is expired.
+
+        Returns:
+            True if token is expired, False otherwise
+        """
+        return datetime.now() >= self.expiry_time
+
+    def to_dict(self) -> Dict[str, str]:
+        """
+        Convert token to dictionary format.
+
+        Returns:
+            Dictionary with access_token and token_type
+        """
+        return {
+            "access_token": self.access_token,
+            "token_type": self.token_type,
+        }
+
+
+class TokenFederationProvider(AuthProvider):
+    """
+    Implementation of Token Federation for Databricks SQL Python driver.
+
+    This provider exchanges third-party access tokens for Databricks in-house tokens
+    when the token issuer is different from the Databricks host.
+    """
+
+    TOKEN_EXCHANGE_ENDPOINT = "/oidc/v1/token"
+    TOKEN_EXCHANGE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:token-exchange"
+    TOKEN_EXCHANGE_SUBJECT_TYPE = "urn:ietf:params:oauth:token-type:jwt"
+
+    def __init__(
+        self,
+        hostname: str,
+        external_provider: AuthProvider,
+        http_client,
+        identity_federation_client_id: Optional[str] = None,
+    ):
+        """
+        Initialize the Token Federation Provider.
+
+        Args:
+            hostname: The Databricks workspace hostname
+            external_provider: The external authentication provider
+            http_client: HTTP client for making requests (required)
+            identity_federation_client_id: Optional client ID for token federation
+        """
+        if not http_client:
+            raise ValueError("http_client is required for TokenFederationProvider")
+
+        self.hostname = parse_hostname(hostname)
+        self.external_provider = external_provider
+        self.http_client = http_client
+        self.identity_federation_client_id = identity_federation_client_id
+
+        self._cached_token: Optional[Token] = None
+        self._external_headers: Dict[str, str] = {}
+
+    def add_headers(self, request_headers: Dict[str, str]):
+        """Add authentication headers to the request."""
+
+        if self._cached_token and not self._cached_token.is_expired():
+            request_headers[
+                "Authorization"
+            ] = f"{self._cached_token.token_type} {self._cached_token.access_token}"
+            return
+
+        # Get the external headers first to check if we need token federation
+        self._external_headers = {}
+        self.external_provider.add_headers(self._external_headers)
+
+        # If no Authorization header from external provider, pass through all headers
+        if "Authorization" not in self._external_headers:
+            request_headers.update(self._external_headers)
+            return
+
+        token = self._get_token()
+        request_headers["Authorization"] = f"{token.token_type} {token.access_token}"
+
+    def _get_token(self) -> Token:
+        """Get or refresh the authentication token."""
+        # Check if cached token is still valid
+        if self._cached_token and not self._cached_token.is_expired():
+            return self._cached_token
+
+        # Extract token from already-fetched headers
+        auth_header = self._external_headers.get("Authorization", "")
+        token_type, access_token = self._extract_token_from_header(auth_header)
+
+        # Check if token exchange is needed
+        if self._should_exchange_token(access_token):
+            try:
+                token = self._exchange_token(access_token)
+                self._cached_token = token
+                return token
+            except Exception as e:
+                logger.warning("Token exchange failed, using external token: %s", e)
+
+        # Use external token directly
+        token = Token(access_token, token_type)
+        self._cached_token = token
+        return token
+
+    def _should_exchange_token(self, access_token: str) -> bool:
+        """Check if the token should be exchanged based on issuer."""
+        decoded = decode_token(access_token)
+        if not decoded:
+            return False
+
+        issuer = decoded.get("iss", "")
+        # Check if issuer host is different from Databricks host
+        return not is_same_host(issuer, self.hostname)
+
+    def _exchange_token(self, access_token: str) -> Token:
+        """Exchange the external token for a Databricks token."""
+        token_url = f"{self.hostname.rstrip('/')}{self.TOKEN_EXCHANGE_ENDPOINT}"
+
+        data = {
+            "grant_type": self.TOKEN_EXCHANGE_GRANT_TYPE,
+            "subject_token": access_token,
+            "subject_token_type": self.TOKEN_EXCHANGE_SUBJECT_TYPE,
+            "scope": "sql",
+            "return_original_token_if_authenticated": "true",
+        }
+
+        if self.identity_federation_client_id:
+            data["client_id"] = self.identity_federation_client_id
+
+        headers = {
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "*/*",
+        }
+
+        body = urlencode(data)
+
+        response = self.http_client.request(
+            HttpMethod.POST, url=token_url, body=body, headers=headers
+        )
+
+        token_response = json.loads(response.data.decode())
+
+        return Token(
+            token_response["access_token"], token_response.get("token_type", "Bearer")
+        )
+
+    def _extract_token_from_header(self, auth_header: str) -> Tuple[str, str]:
+        """Extract token type and access token from Authorization header."""
+        if not auth_header:
+            raise ValueError("Authorization header is missing")
+
+        parts = auth_header.split(" ", 1)
+        if len(parts) != 2:
+            raise ValueError("Invalid Authorization header format")
+
+        return parts[0], parts[1]

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/backend/sea/utils/http_client.py

@@ -90,6 +90,9 @@ class SeaHttpClient:
         )
         self._retry_delay_default = kwargs.get("_retry_delay_default", 5.0)
         self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
+        self._respect_server_retry_after_header = kwargs.get(
+            "_respect_server_retry_after_header", False
+        )
 
         # Connection pooling settings
         self.max_connections = kwargs.get("max_connections", 10)
@@ -114,6 +117,7 @@ class SeaHttpClient:
                 stop_after_attempts_duration=self._retry_stop_after_attempts_duration,
                 delay_default=self._retry_delay_default,
                 force_dangerous_codes=self.force_dangerous_codes,
+                respect_server_retry_after_header=self._respect_server_retry_after_header,
                 urllib3_kwargs=urllib3_kwargs,
             )
         else:

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/backend/thrift_backend.py

@@ -189,6 +189,9 @@ class ThriftDatabricksClient(DatabricksClient):
                 " This behaviour is deprecated and will be removed in a future release."
             )
         self.force_dangerous_codes = kwargs.get("_retry_dangerous_codes", [])
+        self._respect_server_retry_after_header = kwargs.get(
+            "_respect_server_retry_after_header", False
+        )
 
         additional_transport_args = {}
 
@@ -215,6 +218,7 @@ class ThriftDatabricksClient(DatabricksClient):
                 stop_after_attempts_duration=self._retry_stop_after_attempts_duration,
                 delay_default=self._retry_delay_default,
                 force_dangerous_codes=self.force_dangerous_codes,
+                respect_server_retry_after_header=self._respect_server_retry_after_header,
                 urllib3_kwargs=urllib3_kwargs,
             )
 

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/client.py

@@ -200,6 +200,12 @@ class Connection:
                     STRUCT is returned as Dict[str, Any]
                     ARRAY is returned as numpy.ndarray
                 When False, complex types are returned as a strings. These are generally deserializable as JSON.
+            :param enable_metric_view_metadata: `bool`, optional (default is False)
+                When True, enables metric view metadata support by setting the
+                spark.sql.thriftserver.metadata.metricview.enabled session configuration.
+                This allows
+                1. cursor.tables() to return METRIC_VIEW table type
+                2. cursor.columns() to return "measure" column type
         """
 
         # Internal arguments in **kwargs:
@@ -248,6 +254,14 @@ class Connection:
             access_token_kv = {"access_token": access_token}
             kwargs = {**kwargs, **access_token_kv}
 
+        enable_metric_view_metadata = kwargs.get("enable_metric_view_metadata", False)
+        if enable_metric_view_metadata:
+            if session_configuration is None:
+                session_configuration = {}
+            session_configuration[
+                "spark.sql.thriftserver.metadata.metricview.enabled"
+            ] = "true"
+
         self.disable_pandas = kwargs.get("_disable_pandas", False)
         self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/common/unified_http_client.py

@@ -99,6 +99,7 @@ class UnifiedHttpClient:
             stop_after_attempts_duration=self.config.retry_stop_after_attempts_duration,
             delay_default=self.config.retry_delay_default,
             force_dangerous_codes=self.config.retry_dangerous_codes,
+            respect_server_retry_after_header=self.config.respect_server_retry_after_header,
         )
 
         # Initialize the required attributes that DatabricksRetryPolicy expects

{databricks_sql_connector-4.1.3 → databricks_sql_connector-4.1.5}/src/databricks/sql/utils.py

@@ -919,6 +919,9 @@ 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"),
+        respect_server_retry_after_header=kwargs.get(
+            "_respect_server_retry_after_header"
+        ),
         proxy_auth_method=kwargs.get("_proxy_auth_method"),
         pool_connections=kwargs.get("_pool_connections"),
         pool_maxsize=kwargs.get("_pool_maxsize"),