basic setup

Author: jprakash-db

poetry.lock

@@ -20,16 +20,19 @@ 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 = { version = "^2.0.0", optional = true }
+
 
 [tool.poetry.extras]
 pyarrow = ["pyarrow"]
+jwt = ["pyjwt"]
 
-[tool.poetry.dev-dependencies]
+[tool.poetry.group.dev.dependencies]
 pytest = "^7.1.2"
 mypy = "^1.10.1"
 pylint = ">=2.12.0"

pyproject.toml

@@ -1,15 +1,19 @@
 import abc
-import base64
+import jwt
 import logging
+import time
 from typing import Callable, Dict, List
-
+from databricks.sql.common.http import HttpMethod, DatabricksHttpClient, HttpHeader
 from databricks.sql.auth.oauth import OAuthManager
-from databricks.sql.auth.endpoint import get_oauth_endpoints, infer_cloud_from_host
+from databricks.sql.auth.endpoint import get_oauth_endpoints
+from databricks.sql.common.http import DatabricksHttpClient, OAuthResponse
 
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
 from databricks.sql.experimental.oauth_persistence import OAuthToken, OAuthPersistence
 
+logger = logging.getLogger(__name__)
+
 
 class AuthProvider:
     def add_headers(self, request_headers: Dict[str, str]):
@@ -33,6 +37,35 @@ def __call__(self, *args, **kwargs) -> HeaderFactory:
         ...
 
 
+class Token:
+    """
+    A class to represent a token.
+
+    Attributes:
+        access_token (str): The access token string.
+        token_type (str): The type of token (e.g., "Bearer").
+        refresh_token (str): The refresh token string.
+    """
+
+    def __init__(self, access_token: str, token_type: str, refresh_token: str):
+        self.access_token = access_token
+        self.token_type = token_type
+        self.refresh_token = refresh_token
+
+    def is_expired(self):
+        try:
+            decoded_token = jwt.decode(
+                self.access_token, options={"verify_signature": False}
+            )
+            exp_time = decoded_token.get("exp")
+            current_time = time.time()
+            buffer_time = 30  # 30 seconds buffer
+            return exp_time and (exp_time - buffer_time) <= current_time
+        except Exception as e:
+            logger.error("Failed to decode token: %s", e)
+            return e
+
+
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
 class AccessTokenAuthProvider(AuthProvider):
@@ -146,3 +179,74 @@ 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:
+        client_id (str): The Azure service principal's client ID.
+        client_secret (str): The Azure service principal's client secret.
+        tenant_id (str): The Azure AD tenant ID.
+    """
+
+    AZURE_AAD_ENDPOINT = "https://login.microsoftonline.com"
+    DATABRICKS_SCOPE = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default"
+    AZURE_TOKEN_ENDPOINT = "oauth2/token"
+
+    def __init__(self, client_id: str, client_secret: str, tenant_id: str):
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.tenant_id = tenant_id
+        self._token: Token = None
+        self._http_client = DatabricksHttpClient.get_instance()
+
+    def auth_type(self) -> str:
+        return "azure-service-principal"
+
+    def __call__(self, *args, **kwargs) -> HeaderFactory:
+        def header_factory() -> Dict[str, str]:
+            self._refresh()
+            return {HttpHeader.AUTHORIZATION: f"Bearer {self._token.access_token}"}
+
+        return header_factory
+
+    def _refresh(self) -> None:
+        if self._token is None or self._token.is_expired():
+            self._token = self._get_token()
+
+    def _get_token(self) -> Token:
+        request_url = (
+            f"{self.AZURE_AAD_ENDPOINT}/{self.tenant_id}/{self.AZURE_TOKEN_ENDPOINT}"
+        )
+        headers = {
+            HttpHeader.CONTENT_TYPE: "application/x-www-form-urlencoded",
+        }
+        data = {
+            "grant_type": "client_credentials",
+            "client_id": self.client_id,
+            "client_secret": self.client_secret,
+            "scope": self.DATABRICKS_SCOPE,
+        }
+
+        response = self._http_client.execute(
+            method=HttpMethod.POST, url=request_url, headers=headers, data=data
+        )
+
+        if response.status_code == 200:
+            oauth_response = OAuthResponse(**response.json())
+            return Token(
+                oauth_response.access_token,
+                oauth_response.token_type,
+                oauth_response.refresh_token,
+            )
+        else:
+            raise Exception(
+                f"Failed to get token: {response.status_code} {response.text}"
+            )

src/databricks/sql/auth/authenticators.py

@@ -0,0 +1,64 @@
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from enum import Enum
+import threading
+from dataclasses import dataclass
+
+# Enums for HTTP Methods
+class HttpMethod(str, Enum):
+    GET = "GET"
+    POST = "POST"
+    PUT = "PUT"
+    DELETE = "DELETE"
+
+
+class HttpHeader(str, Enum):
+    CONTENT_TYPE = "Content-Type"
+    AUTHORIZATION = "Authorization"
+
+
+# Dataclass for HTTP Response
+@dataclass
+class OAuthResponse:
+    token_type: str = ""
+    expires_in: int = 0
+    ext_expires_in: int = 0
+    expires_on: int = 0
+    not_before: int = 0
+    resource: str = ""
+    access_token: str = ""
+    refresh_token: str = ""
+
+
+# Singleton class for common Http Client
+class DatabricksHttpClient:
+    ## TODO: Unify all the http clients in the PySQL Connector
+
+    _instance = None
+    _lock = threading.Lock()
+
+    def __init__(self):
+        self.session = requests.Session()
+        adapter = HTTPAdapter(
+            pool_connections=5,
+            pool_maxsize=10,
+            max_retries=Retry(total=10, backoff_factor=0.1),
+        )
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+
+    @classmethod
+    def get_instance(cls) -> "DatabricksHttpClient":
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = DatabricksHttpClient()
+        return cls._instance
+
+    def execute(self, method: HttpMethod, url: str, **kwargs) -> requests.Response:
+        with self.session.request(method, url, **kwargs) as response:
+            return response
+
+    def close(self):
+        self.session.close()

src/databricks/sql/common/http.py

@@ -1,20 +1,26 @@
 import unittest
 import pytest
 from typing import Optional
-from unittest.mock import patch
-
+from unittest.mock import patch, MagicMock
+import jwt
 from databricks.sql.auth.auth import (
     AccessTokenAuthProvider,
     AuthProvider,
     ExternalAuthProvider,
     AuthType,
 )
+import time
+from datetime import datetime, timedelta
 from databricks.sql.auth.auth import (
     get_python_sql_connector_auth_provider,
     PYSQL_OAUTH_CLIENT_ID,
 )
 from databricks.sql.auth.oauth import OAuthManager
-from databricks.sql.auth.authenticators import DatabricksOAuthProvider
+from databricks.sql.auth.authenticators import (
+    DatabricksOAuthProvider,
+    AzureServicePrincipalCredentialProvider,
+    Token,
+)
 from databricks.sql.auth.endpoint import (
     CloudType,
     InHouseOAuthEndpointCollection,
@@ -190,3 +196,95 @@ def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):
         auth_provider = get_python_sql_connector_auth_provider(hostname)
         self.assertTrue(type(auth_provider).__name__, "DatabricksOAuthProvider")
         self.assertTrue(auth_provider._client_id, PYSQL_OAUTH_CLIENT_ID)
+
+
+class TestAzureServicePrincipalCredentialProvider:
+    @pytest.fixture
+    def indefinite_token(self):
+        secret_key = "mysecret"
+        expires_in_100_years = int(time.time()) + (100 * 365 * 24 * 60 * 60)
+
+        payload = {"sub": "user123", "role": "admin", "exp": expires_in_100_years}
+
+        token = jwt.encode(payload, secret_key, algorithm="HS256")
+        return Token(token, "Bearer", "refresh_token")
+
+    @pytest.fixture
+    def http_response(self):
+        def status_response(response_status_code):
+            mock_response = MagicMock()
+            mock_response.status_code = response_status_code
+            mock_response.json.return_value = {
+                "access_token": "abc123",
+                "token_type": "Bearer",
+                "refresh_token": None,
+            }
+            return mock_response
+
+        return status_response
+
+    @pytest.fixture
+    def provider(self):
+        return AzureServicePrincipalCredentialProvider(
+            client_id="dummy-client",
+            client_secret="dummy-secret",
+            tenant_id="dummy-tenant",
+        )
+
+    def test_token_refresh(self, provider):
+        with patch.object(provider, "_get_token") as mock_get_token:
+            mock_get_token.return_value = Token(
+                "access_token", "Bearer", "refresh_token"
+            )
+            header_factory = provider()
+            headers = header_factory()
+
+            assert headers["Authorization"] == "Bearer access_token"
+            mock_get_token.assert_called_once()
+
+    def test_no_token_refresh__when_token_is_not_expired(
+        self, provider, indefinite_token
+    ):
+        with patch.object(provider, "_get_token") as mock_get_token:
+            mock_get_token.return_value = indefinite_token
+
+            # Call the provider multiple times
+            header_factory1 = provider()
+            header_factory2 = provider()
+            header_factory3 = provider()
+
+            # Get headers from each factory
+            headers1 = header_factory1()
+            headers2 = header_factory2()
+            headers3 = header_factory3()
+
+            # Verify _get_token was called only once
+            mock_get_token.assert_called_once()
+
+            # Verify all headers contain the same token
+            expected_auth_header = f"Bearer {indefinite_token.access_token}"
+            assert headers1["Authorization"] == expected_auth_header
+            assert headers2["Authorization"] == expected_auth_header
+            assert headers3["Authorization"] == expected_auth_header
+
+    def test_get_token_success(self, provider, http_response):
+
+        # Patch the HTTP client's execute method
+        with patch.object(
+            provider._http_client, "execute", return_value=http_response(200)
+        ) as mock_execute:
+            token = provider._get_token()
+
+            # Assert
+            assert isinstance(token, Token)
+            assert token.access_token == "abc123"
+            assert token.token_type == "Bearer"
+            assert token.refresh_token is None
+
+    def test_get_token_failure(self, provider, http_response):
+        with patch.object(
+            provider._http_client, "execute", return_value=http_response(400)
+        ) as mock_execute:
+            with pytest.raises(Exception) as e:
+                provider._get_token()
+            assert "Failed to get token: 400" in str(e.value)