Merge branch 'main' into enhance-ci

Author: varun-edachali-dbx

.github/CODEOWNERS

@@ -1,5 +1,8 @@
 # Release History
 
+# 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)
+
 # 4.0.4 (2025-06-16)
 
 - Update thrift client library after cleaning up unused fields and structs (databricks/databricks-sql-python#553 by @vikrantpuppala)

CHANGELOG.md

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "4.0.4"
+version = "4.0.5"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"
@@ -20,16 +20,18 @@ 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"
+
 
 [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"

poetry.lock

@@ -68,7 +68,7 @@ def __repr__(self):
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "4.0.4"
+__version__ = "4.0.5"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

pyproject.toml

@@ -1,51 +1,29 @@
-from enum import Enum
 from typing import Optional, List
 
 from databricks.sql.auth.authenticators import (
     AuthProvider,
     AccessTokenAuthProvider,
     ExternalAuthProvider,
     DatabricksOAuthProvider,
+    AzureServicePrincipalCredentialProvider,
 )
-
-
-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
+from databricks.sql.auth.common import AuthType, ClientContext
 
 
 def get_auth_provider(cfg: ClientContext):
     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,
+                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
@@ -102,10 +80,13 @@ def get_client_id_and_redirect_port(use_azure_auth: bool):
 
 
 def get_python_sql_connector_auth_provider(hostname: str, **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,6 +101,10 @@ 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,

src/databricks/sql/__init__.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.
@@ -146,3 +154,82 @@ def add_headers(self, request_headers: Dict[str, str]):
         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,
+        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
+        )
+
+    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,
+            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