testing sdk

Author: jprakash-db

src/databricks/sql/auth/auth.py

@@ -116,9 +116,7 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
     client_id = kwargs.get("oauth_client_id")
     redirect_port_range = kwargs.get("oauth_redirect_port_range")
 
-    if auth_type == AuthType.AZURE_SP_M2M.value:
-        pass
-    else:
+    if auth_type != AuthType.AZURE_SP_M2M.value:
         (client_id, redirect_port_range) = get_client_id_and_redirect_port(
             auth_type == AuthType.AZURE_OAUTH.value
         )
@@ -140,7 +138,7 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
         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 client_id and kwargs.get("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"),

src/databricks/sql/auth/authenticators.py

@@ -1,6 +1,6 @@
 import abc
 import logging
-from typing import Callable, Dict, List, Optional
+from typing import Callable, Dict, List
 from databricks.sql.common.http import HttpHeader
 from databricks.sql.auth.oauth import (
     OAuthManager,
@@ -9,12 +9,12 @@
 )
 from databricks.sql.auth.endpoint import get_oauth_endpoints
 from databricks.sql.auth.common import AuthType, get_effective_azure_login_app_id
+from databricks.sdk import WorkspaceClient
 
 # 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:
@@ -189,6 +189,13 @@ def __init__(
         azure_tenant_id,
         azure_workspace_resource_id=None,
     ):
+        self.workspace_api_client = WorkspaceClient(
+            host=hostname,
+            azure_workspace_resource_id=azure_workspace_resource_id,
+            azure_tenant_id=azure_tenant_id,
+            azure_client_id=oauth_client_id,
+            azure_client_secret=oauth_client_secret,
+        )
         self.hostname = hostname
         self.oauth_client_id = oauth_client_id
         self.oauth_client_secret = oauth_client_secret
@@ -207,25 +214,26 @@ def get_token_source(self, resource: str) -> RefreshableTokenSource:
         )
 
     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)
+        # 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()
+            # 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,
-            }
+            # 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
+            # if self.azure_workspace_resource_id:
+            #     headers[
+            #         self.DATABRICKS_AZURE_WORKSPACE_RESOURCE_ID_HEADER
+            #     ] = self.azure_workspace_resource_id
 
-            return headers
+            # return headers
+            return self.workspace_api_client.config.authenticate()
 
         return header_factory

src/databricks/sql/auth/common.py

@@ -1,5 +1,4 @@
 from enum import Enum
-from typing import Optional
 
 
 class AuthType(Enum):
@@ -8,21 +7,23 @@ class AuthType(Enum):
     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")
+
+
 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.
 
     """
-    azure_app_ids = {
-        ".dev.azuredatabricks.net": "62a912ac-b58e-4c1d-89ea-b2dbfc7358fc",
-        ".staging.azuredatabricks.net": "4a67d088-db5c-48f1-9ff2-0aace800ae68",
-    }
-
-    for domain, app_id in azure_app_ids.items():
+    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 "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+    return AzureAppId.PROD.value

src/databricks/sql/auth/oauth.py

@@ -39,7 +39,7 @@ def __init__(self, access_token: str, token_type: str, refresh_token: str):
         self.token_type = token_type
         self.refresh_token = refresh_token
 
-    def is_expired(self):
+    def is_expired(self) -> bool:
         try:
             decoded_token = jwt.decode(
                 self.access_token, options={"verify_signature": False}
@@ -50,7 +50,7 @@ def is_expired(self):
             return exp_time and (exp_time - buffer_time) <= current_time
         except Exception as e:
             logger.error("Failed to decode token: %s", e)
-            return e
+            raise e
 
 
 class RefreshableTokenSource(ABC):
@@ -347,18 +347,17 @@ def refresh(self) -> Token:
             }
         )
 
-        response = self._http_client.execute(
+        with self._http_client.execute(
             method=HttpMethod.POST, url=self.token_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}"
-            )
+        ) as response:
+            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/common/http.py

@@ -4,6 +4,12 @@
 from enum import Enum
 import threading
 from dataclasses import dataclass
+from contextlib import contextmanager
+from typing import Generator
+import logging
+
+logger = logging.getLogger(__name__)
+
 
 # Enums for HTTP Methods
 class HttpMethod(str, Enum):
@@ -56,10 +62,19 @@ def get_instance(cls) -> "DatabricksHttpClient":
                 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.value, url, **kwargs) as response:
-            return response
+    
+    @contextmanager
+    def execute(self, method: HttpMethod, url: str, **kwargs) -> Generator[requests.Response, None, None]:
+        response = None
+        try:
+            response = self.session.request(method.value, url, **kwargs)
+            yield response
+        except Exception as e:
+            logger.error(f"Error executing HTTP request in DatabricksHttpClient: {e}")
+            raise e
+        finally:
+            if response is not None:
+                response.close()
 
     def close(self):
         self.session.close()

tests/unit/test_auth.py

@@ -1,6 +1,5 @@
 import unittest
 import pytest
-from typing import Optional
 from unittest.mock import patch, MagicMock
 import jwt
 from databricks.sql.auth.auth import (
@@ -10,7 +9,6 @@
     AuthType,
 )
 import time
-from datetime import datetime, timedelta
 from databricks.sql.auth.auth import (
     get_python_sql_connector_auth_provider,
     PYSQL_OAUTH_CLIENT_ID,