databricks-sql-connector 4.0.5__tar.gz → 4.1.0__tar.gz

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/CHANGELOG.md

@@ -1,5 +1,29 @@
 # Release History
 
+# 4.1.0 (2025-08-18)
+- Removed Codeowners (databricks/databricks-sql-python#623 by @jprakash-db)
+- Azure Service Principal Credential Provider (databricks/databricks-sql-python#621 by @jprakash-db)
+- Add optional telemetry support to the python connector (databricks/databricks-sql-python#628 by @saishreeeee)
+- Fix potential resource leak in `CloudFetchQueue` (databricks/databricks-sql-python#624 by @varun-edachali-dbx)
+- Generalise Backend Layer (databricks/databricks-sql-python#604 by @varun-edachali-dbx)
+- Arrow performance optimizations (databricks/databricks-sql-python#638 by @jprakash-db)
+- Connection errors to unauthenticated telemetry endpoint (databricks/databricks-sql-python#619 by @saishreeeee)
+- SEA: Execution Phase (databricks/databricks-sql-python#645 by @varun-edachali-dbx)
+- Add retry mechanism to telemetry requests (databricks/databricks-sql-python#617 by @saishreeeee)
+- SEA: Fetch Phase (databricks/databricks-sql-python#650 by @varun-edachali-dbx)
+- added logs for cloud fetch speed (databricks/databricks-sql-python#654 by @shivam2680)
+- Make telemetry batch size configurable and add time-based flush (databricks/databricks-sql-python#622 by @saishreeeee)
+- Normalise type code (databricks/databricks-sql-python#652 by @varun-edachali-dbx)
+- Testing for telemetry (databricks/databricks-sql-python#616 by @saishreeeee)
+- Bug fixes in telemetry (databricks/databricks-sql-python#659 by @saishreeeee)
+- Telemetry server-side flag integration (databricks/databricks-sql-python#646 by @saishreeeee)
+- Enhance SEA HTTP Client (databricks/databricks-sql-python#618 by @varun-edachali-dbx)
+- SEA: Allow large metadata responses (databricks/databricks-sql-python#653 by @varun-edachali-dbx)
+- Added code coverage workflow to test the code coverage from unit and e2e tests (databricks/databricks-sql-python#657 by @msrathore-db)
+- Concat tables to be backward compatible (databricks/databricks-sql-python#647 by @jprakash-db)
+- Refactor codebase to use a unified http client (databricks/databricks-sql-python#673 by @vikrantpuppala)
+- Add kerberos support for proxy auth (databricks/databricks-sql-python#675 by @vikrantpuppala)
+
 # 4.0.5 (2025-06-24)
 - Fix: Reverted change in cursor close handling which led to errors impacting users (databricks/databricks-sql-python#613 by @madhav-db)
 

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: databricks-sql-connector
-Version: 4.0.5
+Version: 4.1.0
 Summary: Databricks SQL Connector for Python
 License: Apache-2.0
 Author: Databricks
@@ -22,6 +22,7 @@ Requires-Dist: pandas (>=1.2.5,<2.3.0) ; python_version >= "3.8" and python_vers
 Requires-Dist: pandas (>=2.2.3,<2.3.0) ; python_version >= "3.13"
 Requires-Dist: pyarrow (>=14.0.1) ; (python_version >= "3.8" and python_version < "3.13") and (extra == "pyarrow")
 Requires-Dist: pyarrow (>=18.0.0) ; (python_version >= "3.13") and (extra == "pyarrow")
+Requires-Dist: pyjwt (>=2.0.0,<3.0.0)
 Requires-Dist: python-dateutil (>=2.8.0,<3.0.0)
 Requires-Dist: requests (>=2.18.1,<3.0.0)
 Requires-Dist: thrift (>=0.16.0,<0.21.0)

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.0.5"
+version = "4.1.0"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"
@@ -20,21 +20,25 @@ requests = "^2.18.1"
 oauthlib = "^3.1.0"
 openpyxl = "^3.0.10"
 urllib3 = ">=1.26"
+python-dateutil = "^2.8.0"
 pyarrow = [
     { version = ">=14.0.1", python = ">=3.8,<3.13", optional=true },
     { version = ">=18.0.0", python = ">=3.13", optional=true }
 ]
-python-dateutil = "^2.8.0"
+pyjwt = "^2.0.0"
+requests-kerberos = {version = "^0.15.0", optional = true}
+
 
 [tool.poetry.extras]
 pyarrow = ["pyarrow"]
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 pytest = "^7.1.2"
 mypy = "^1.10.1"
 pylint = ">=2.12.0"
 black = "^22.3.0"
 pytest-dotenv = "^0.5.2"
+pytest-cov = "^4.0.0"
 numpy = [
     { version = ">=1.16.6", python = ">=3.8,<3.11" },
     { version = ">=1.23.4", python = ">=3.11" },
@@ -62,3 +66,21 @@ log_cli = "false"
 log_cli_level = "INFO"
 testpaths = ["tests"]
 env_files = ["test.env"]
+
+[tool.coverage.run]
+source = ["src"]
+branch = true
+omit = [
+    "*/tests/*",
+    "*/test_*",
+    "*/__pycache__/*",
+    "*/thrift_api/*",  
+]
+
+[tool.coverage.report]
+precision = 2
+show_missing = true
+skip_covered = false
+
+[tool.coverage.xml]
+output = "coverage.xml"

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/src/databricks/sql/__init__.py

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

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/src/databricks/sql/auth/auth.py

@@ -1,4 +1,3 @@
-from enum import Enum
 from typing import Optional, List
 
 from databricks.sql.auth.authenticators import (
@@ -6,46 +5,26 @@ from databricks.sql.auth.authenticators import (
     AccessTokenAuthProvider,
     ExternalAuthProvider,
     DatabricksOAuthProvider,
+    AzureServicePrincipalCredentialProvider,
 )
+from databricks.sql.auth.common import AuthType, ClientContext
 
 
-class AuthType(Enum):
-    DATABRICKS_OAUTH = "databricks-oauth"
-    AZURE_OAUTH = "azure-oauth"
-    # other supported types (access_token) can be inferred
-    # we can add more types as needed later
-
-
-class ClientContext:
-    def __init__(
-        self,
-        hostname: str,
-        access_token: Optional[str] = None,
-        auth_type: Optional[str] = None,
-        oauth_scopes: Optional[List[str]] = None,
-        oauth_client_id: Optional[str] = None,
-        oauth_redirect_port_range: Optional[List[int]] = None,
-        use_cert_as_auth: Optional[str] = None,
-        tls_client_cert_file: Optional[str] = None,
-        oauth_persistence=None,
-        credentials_provider=None,
-    ):
-        self.hostname = hostname
-        self.access_token = access_token
-        self.auth_type = auth_type
-        self.oauth_scopes = oauth_scopes
-        self.oauth_client_id = oauth_client_id
-        self.oauth_redirect_port_range = oauth_redirect_port_range
-        self.use_cert_as_auth = use_cert_as_auth
-        self.tls_client_cert_file = tls_client_cert_file
-        self.oauth_persistence = oauth_persistence
-        self.credentials_provider = credentials_provider
-
-
-def get_auth_provider(cfg: ClientContext):
+def get_auth_provider(cfg: ClientContext, http_client):
     if cfg.credentials_provider:
         return ExternalAuthProvider(cfg.credentials_provider)
-    if cfg.auth_type in [AuthType.DATABRICKS_OAUTH.value, AuthType.AZURE_OAUTH.value]:
+    elif cfg.auth_type == AuthType.AZURE_SP_M2M.value:
+        return ExternalAuthProvider(
+            AzureServicePrincipalCredentialProvider(
+                cfg.hostname,
+                cfg.azure_client_id,
+                cfg.azure_client_secret,
+                http_client,
+                cfg.azure_tenant_id,
+                cfg.azure_workspace_resource_id,
+            )
+        )
+    elif cfg.auth_type in [AuthType.DATABRICKS_OAUTH.value, AuthType.AZURE_OAUTH.value]:
         assert cfg.oauth_redirect_port_range is not None
         assert cfg.oauth_client_id is not None
         assert cfg.oauth_scopes is not None
@@ -56,6 +35,7 @@ def get_auth_provider(cfg: ClientContext):
             cfg.oauth_redirect_port_range,
             cfg.oauth_client_id,
             cfg.oauth_scopes,
+            http_client,
             cfg.auth_type,
         )
     elif cfg.access_token is not None:
@@ -75,6 +55,8 @@ def get_auth_provider(cfg: ClientContext):
                 cfg.oauth_redirect_port_range,
                 cfg.oauth_client_id,
                 cfg.oauth_scopes,
+                http_client,
+                cfg.auth_type or AuthType.DATABRICKS_OAUTH.value,
             )
         else:
             raise RuntimeError("No valid authentication settings!")
@@ -101,11 +83,14 @@ def get_client_id_and_redirect_port(use_azure_auth: bool):
     )
 
 
-def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
+def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs):
+    # TODO : unify all the auth mechanisms with the Python SDK
+
     auth_type = kwargs.get("auth_type")
     (client_id, redirect_port_range) = get_client_id_and_redirect_port(
         auth_type == AuthType.AZURE_OAUTH.value
     )
+
     if kwargs.get("username") or kwargs.get("password"):
         raise ValueError(
             "Username/password authentication is no longer supported. "
@@ -120,10 +105,14 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
         tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
         oauth_scopes=PYSQL_OAUTH_SCOPES,
         oauth_client_id=kwargs.get("oauth_client_id") or client_id,
+        azure_client_id=kwargs.get("azure_client_id"),
+        azure_client_secret=kwargs.get("azure_client_secret"),
+        azure_tenant_id=kwargs.get("azure_tenant_id"),
+        azure_workspace_resource_id=kwargs.get("azure_workspace_resource_id"),
         oauth_redirect_port_range=[kwargs["oauth_redirect_port"]]
         if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port")
         else redirect_port_range,
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
     )
-    return get_auth_provider(cfg)
+    return get_auth_provider(cfg, http_client)

{databricks_sql_connector-4.0.5 → databricks_sql_connector-4.1.0}/src/databricks/sql/auth/authenticators.py

@@ -1,10 +1,18 @@
 import abc
-import base64
 import logging
 from typing import Callable, Dict, List
-
-from databricks.sql.auth.oauth import OAuthManager
-from databricks.sql.auth.endpoint import get_oauth_endpoints, infer_cloud_from_host
+from databricks.sql.common.http import HttpHeader
+from databricks.sql.auth.oauth import (
+    OAuthManager,
+    RefreshableTokenSource,
+    ClientCredentialsTokenSource,
+)
+from databricks.sql.auth.endpoint import get_oauth_endpoints
+from databricks.sql.auth.common import (
+    AuthType,
+    get_effective_azure_login_app_id,
+    get_azure_tenant_id_from_host,
+)
 
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
@@ -55,6 +63,7 @@ class DatabricksOAuthProvider(AuthProvider):
         redirect_port_range: List[int],
         client_id: str,
         scopes: List[str],
+        http_client,
         auth_type: str = "databricks-oauth",
     ):
         try:
@@ -71,6 +80,7 @@ class DatabricksOAuthProvider(AuthProvider):
                 port_range=redirect_port_range,
                 client_id=client_id,
                 idp_endpoint=idp_endpoint,
+                http_client=http_client,
             )
             self._hostname = hostname
             self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(cloud_scopes)
@@ -146,3 +156,85 @@ class ExternalAuthProvider(AuthProvider):
         headers = self._header_factory()
         for k, v in headers.items():
             request_headers[k] = v
+
+
+class AzureServicePrincipalCredentialProvider(CredentialsProvider):
+    """
+    A credential provider for Azure Service Principal authentication with Databricks.
+
+    This class implements the CredentialsProvider protocol to authenticate requests
+    to Databricks REST APIs using Azure Active Directory (AAD) service principal
+    credentials. It handles OAuth 2.0 client credentials flow to obtain access tokens
+    from Azure AD and automatically refreshes them when they expire.
+
+    Attributes:
+        hostname (str): The Databricks workspace hostname.
+        azure_client_id (str): The Azure service principal's client ID.
+        azure_client_secret (str): The Azure service principal's client secret.
+        azure_tenant_id (str): The Azure AD tenant ID.
+        azure_workspace_resource_id (str, optional): The Azure workspace resource ID.
+    """
+
+    AZURE_AAD_ENDPOINT = "https://login.microsoftonline.com"
+    AZURE_TOKEN_ENDPOINT = "oauth2/token"
+
+    AZURE_MANAGED_RESOURCE = "https://management.core.windows.net/"
+
+    DATABRICKS_AZURE_SP_TOKEN_HEADER = "X-Databricks-Azure-SP-Management-Token"
+    DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID_HEADER = (
+        "X-Databricks-Azure-Workspace-Resource-Id"
+    )
+
+    def __init__(
+        self,
+        hostname,
+        azure_client_id,
+        azure_client_secret,
+        http_client,
+        azure_tenant_id=None,
+        azure_workspace_resource_id=None,
+    ):
+        self.hostname = hostname
+        self.azure_client_id = azure_client_id
+        self.azure_client_secret = azure_client_secret
+        self.azure_workspace_resource_id = azure_workspace_resource_id
+        self.azure_tenant_id = azure_tenant_id or get_azure_tenant_id_from_host(
+            hostname, http_client
+        )
+        self._http_client = http_client
+
+    def auth_type(self) -> str:
+        return AuthType.AZURE_SP_M2M.value
+
+    def get_token_source(self, resource: str) -> RefreshableTokenSource:
+        return ClientCredentialsTokenSource(
+            token_url=f"{self.AZURE_AAD_ENDPOINT}/{self.azure_tenant_id}/{self.AZURE_TOKEN_ENDPOINT}",
+            client_id=self.azure_client_id,
+            client_secret=self.azure_client_secret,
+            http_client=self._http_client,
+            extra_params={"resource": resource},
+        )
+
+    def __call__(self, *args, **kwargs) -> HeaderFactory:
+        inner = self.get_token_source(
+            resource=get_effective_azure_login_app_id(self.hostname)
+        )
+        cloud = self.get_token_source(resource=self.AZURE_MANAGED_RESOURCE)
+
+        def header_factory() -> Dict[str, str]:
+            inner_token = inner.get_token()
+            cloud_token = cloud.get_token()
+
+            headers = {
+                HttpHeader.AUTHORIZATION.value: f"{inner_token.token_type} {inner_token.access_token}",
+                self.DATABRICKS_AZURE_SP_TOKEN_HEADER: cloud_token.access_token,
+            }
+
+            if self.azure_workspace_resource_id:
+                headers[
+                    self.DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID_HEADER
+                ] = self.azure_workspace_resource_id
+
+            return headers
+
+        return header_factory

databricks_sql_connector-4.1.0/src/databricks/sql/auth/common.py

@@ -0,0 +1,127 @@
+from enum import Enum
+import logging
+from typing import Optional, List
+from urllib.parse import urlparse
+from databricks.sql.auth.retry import DatabricksRetryPolicy
+from databricks.sql.common.http import HttpMethod
+
+logger = logging.getLogger(__name__)
+
+
+class AuthType(Enum):
+    DATABRICKS_OAUTH = "databricks-oauth"
+    AZURE_OAUTH = "azure-oauth"
+    AZURE_SP_M2M = "azure-sp-m2m"
+
+
+class AzureAppId(Enum):
+    DEV = (".dev.azuredatabricks.net", "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc")
+    STAGING = (".staging.azuredatabricks.net", "4a67d088-db5c-48f1-9ff2-0aace800ae68")
+    PROD = (".azuredatabricks.net", "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d")
+
+
+class ClientContext:
+    def __init__(
+        self,
+        hostname: str,
+        access_token: Optional[str] = None,
+        auth_type: Optional[str] = None,
+        oauth_scopes: Optional[List[str]] = None,
+        oauth_client_id: Optional[str] = None,
+        azure_client_id: Optional[str] = None,
+        azure_client_secret: Optional[str] = None,
+        azure_tenant_id: Optional[str] = None,
+        azure_workspace_resource_id: Optional[str] = None,
+        oauth_redirect_port_range: Optional[List[int]] = None,
+        use_cert_as_auth: Optional[str] = None,
+        tls_client_cert_file: Optional[str] = None,
+        oauth_persistence=None,
+        credentials_provider=None,
+        # HTTP client configuration parameters
+        ssl_options=None,  # SSLOptions type
+        socket_timeout: Optional[float] = None,
+        retry_stop_after_attempts_count: Optional[int] = None,
+        retry_delay_min: Optional[float] = None,
+        retry_delay_max: Optional[float] = None,
+        retry_stop_after_attempts_duration: Optional[float] = None,
+        retry_delay_default: Optional[float] = None,
+        retry_dangerous_codes: Optional[List[int]] = None,
+        proxy_auth_method: Optional[str] = None,
+        pool_connections: Optional[int] = None,
+        pool_maxsize: Optional[int] = None,
+        user_agent: Optional[str] = None,
+    ):
+        self.hostname = hostname
+        self.access_token = access_token
+        self.auth_type = auth_type
+        self.oauth_scopes = oauth_scopes
+        self.oauth_client_id = oauth_client_id
+        self.azure_client_id = azure_client_id
+        self.azure_client_secret = azure_client_secret
+        self.azure_tenant_id = azure_tenant_id
+        self.azure_workspace_resource_id = azure_workspace_resource_id
+        self.oauth_redirect_port_range = oauth_redirect_port_range
+        self.use_cert_as_auth = use_cert_as_auth
+        self.tls_client_cert_file = tls_client_cert_file
+        self.oauth_persistence = oauth_persistence
+        self.credentials_provider = credentials_provider
+
+        # HTTP client configuration
+        self.ssl_options = ssl_options
+        self.socket_timeout = socket_timeout
+        self.retry_stop_after_attempts_count = retry_stop_after_attempts_count or 5
+        self.retry_delay_min = retry_delay_min or 1.0
+        self.retry_delay_max = retry_delay_max or 10.0
+        self.retry_stop_after_attempts_duration = (
+            retry_stop_after_attempts_duration or 300.0
+        )
+        self.retry_delay_default = retry_delay_default or 5.0
+        self.retry_dangerous_codes = retry_dangerous_codes or []
+        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
+
+
+def get_effective_azure_login_app_id(hostname) -> str:
+    """
+    Get the effective Azure login app ID for a given hostname.
+    This function determines the appropriate Azure login app ID based on the hostname.
+    If the hostname does not match any of these domains, it returns the default Databricks resource ID.
+
+    """
+    for azure_app_id in AzureAppId:
+        domain, app_id = azure_app_id.value
+        if domain in hostname:
+            return app_id
+
+    # default databricks resource id
+    return AzureAppId.PROD.value[1]
+
+
+def get_azure_tenant_id_from_host(host: str, http_client) -> str:
+    """
+    Load the Azure tenant ID from the Azure Databricks login page.
+
+    This function retrieves the Azure tenant ID by making a request to the Databricks
+    Azure Active Directory (AAD) authentication endpoint. The endpoint redirects to
+    the Azure login page, and the tenant ID is extracted from the redirect URL.
+    """
+
+    login_url = f"{host}/aad/auth"
+    logger.debug("Loading tenant ID from %s", login_url)
+
+    with http_client.request_context(HttpMethod.GET, login_url) as resp:
+        entra_id_endpoint = resp.retries.history[-1].redirect_location
+        if entra_id_endpoint is None:
+            raise ValueError(
+                f"No Location header in response from {login_url}: {entra_id_endpoint}"
+            )
+
+    # The final redirect URL has the following form: https://login.microsoftonline.com/<tenant-id>/oauth2/authorize?...
+    # The domain may change depending on the Azure cloud (e.g. login.microsoftonline.us for US Government cloud).
+    url = urlparse(entra_id_endpoint)
+    path_segments = url.path.split("/")
+    if len(path_segments) < 2:
+        raise ValueError(f"Invalid path in Location header: {url.path}")
+    return path_segments[1]