[PECO-626] Support OAuth flow for Databricks Azure (databricks#86)

## Summary

Support OAuth flow for Databricks Azure

## Background

Some OAuth endpoints (e.g. Open ID Configuration) and scopes are different between Databricks Azure and AWS. Current code only supports OAuth flow on Databricks in AWS

## What changes are proposed in this pull request?

- Change `OAuthManager` to decouple Databricks AWS specific configuration from OAuth flow
- Add `sql/auth/endpoint.py` that implements cloud specific OAuth endpoint configuration
- Change `DatabricksOAuthProvider` to work with the OAuth configurations in different Databricks cloud (AWS, Azure)
- Add the corresponding unit tests

Author: jackyhu-db

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## 2.6.x (Unreleased)
 
+- Add support for OAuth on Databricks Azure
+
 ## 2.6.2 (2023-06-14)
 
 - Fix: Retry GetOperationStatus requests for http errors

src/databricks/sql/auth/auth.py

@@ -8,6 +8,7 @@
     ExternalAuthProvider,
     DatabricksOAuthProvider,
 )
+from databricks.sql.auth.endpoint import infer_cloud_from_host, CloudType
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 
 
@@ -75,7 +76,9 @@ def get_auth_provider(cfg: ClientContext):
 
 PYSQL_OAUTH_SCOPES = ["sql", "offline_access"]
 PYSQL_OAUTH_CLIENT_ID = "databricks-sql-python"
+PYSQL_OAUTH_AZURE_CLIENT_ID = "96eecda7-19ea-49cc-abb5-240097d554f5"
 PYSQL_OAUTH_REDIRECT_PORT_RANGE = list(range(8020, 8025))
+PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE = [8030]
 
 
 def normalize_host_name(hostname: str):
@@ -84,7 +87,16 @@ def normalize_host_name(hostname: str):
     return f"{maybe_scheme}{hostname}{maybe_trailing_slash}"
 
 
+def get_client_id_and_redirect_port(hostname: str):
+    return (
+        (PYSQL_OAUTH_CLIENT_ID, PYSQL_OAUTH_REDIRECT_PORT_RANGE)
+        if infer_cloud_from_host(hostname) == CloudType.AWS
+        else (PYSQL_OAUTH_AZURE_CLIENT_ID, PYSQL_OAUTH_AZURE_REDIRECT_PORT_RANGE)
+    )
+
+
 def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
+    (client_id, redirect_port_range) = get_client_id_and_redirect_port(hostname)
     cfg = ClientContext(
         hostname=normalize_host_name(hostname),
         auth_type=kwargs.get("auth_type"),
@@ -94,10 +106,10 @@ def get_python_sql_connector_auth_provider(hostname: str, **kwargs):
         use_cert_as_auth=kwargs.get("_use_cert_as_auth"),
         tls_client_cert_file=kwargs.get("_tls_client_cert_file"),
         oauth_scopes=PYSQL_OAUTH_SCOPES,
-        oauth_client_id=kwargs.get("oauth_client_id") or PYSQL_OAUTH_CLIENT_ID,
+        oauth_client_id=kwargs.get("oauth_client_id") or client_id,
         oauth_redirect_port_range=[kwargs["oauth_redirect_port"]]
         if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port")
-        else PYSQL_OAUTH_REDIRECT_PORT_RANGE,
+        else redirect_port_range,
         oauth_persistence=kwargs.get("experimental_oauth_persistence"),
         credentials_provider=kwargs.get("credentials_provider"),
     )

src/databricks/sql/auth/authenticators.py

@@ -4,6 +4,7 @@
 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
 
 # Private API: this is an evolving interface and it will change in the future.
 # Please must not depend on it in your applications.
@@ -70,11 +71,26 @@ def __init__(
         scopes: List[str],
     ):
         try:
+            cloud_type = infer_cloud_from_host(hostname)
+            if not cloud_type:
+                raise NotImplementedError("Cannot infer the cloud type from hostname")
+
+            idp_endpoint = get_oauth_endpoints(cloud_type)
+            if not idp_endpoint:
+                raise NotImplementedError(
+                    f"OAuth is not supported for cloud ${cloud_type.value}"
+                )
+
+            # Convert to the corresponding scopes in the corresponding IdP
+            cloud_scopes = idp_endpoint.get_scopes_mapping(scopes)
+
             self.oauth_manager = OAuthManager(
-                port_range=redirect_port_range, client_id=client_id
+                port_range=redirect_port_range,
+                client_id=client_id,
+                idp_endpoint=idp_endpoint,
             )
             self._hostname = hostname
-            self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(scopes)
+            self._scopes_as_str = DatabricksOAuthProvider.SCOPE_DELIM.join(cloud_scopes)
             self._oauth_persistence = oauth_persistence
             self._client_id = client_id
             self._access_token = None

src/databricks/sql/auth/endpoint.py

@@ -0,0 +1,112 @@
+#
+# It implements all the cloud specific OAuth configuration/metadata
+#
+#   Azure: It uses AAD
+#   AWS: It uses Databricks internal IdP
+#   GCP: Not support yet
+#
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Optional, List
+import os
+
+OIDC_REDIRECTOR_PATH = "oidc"
+
+
+class OAuthScope:
+    OFFLINE_ACCESS = "offline_access"
+    SQL = "sql"
+
+
+class CloudType(Enum):
+    AWS = "aws"
+    AZURE = "azure"
+
+
+DATABRICKS_AWS_DOMAINS = [".cloud.databricks.com", ".dev.databricks.com"]
+DATABRICKS_AZURE_DOMAINS = [
+    ".azuredatabricks.net",
+    ".databricks.azure.cn",
+    ".databricks.azure.us",
+]
+
+
+# Infer cloud type from Databricks SQL instance hostname
+def infer_cloud_from_host(hostname: str) -> Optional[CloudType]:
+    # normalize
+    host = hostname.lower().replace("https://", "").split("/")[0]
+
+    if any(e for e in DATABRICKS_AZURE_DOMAINS if host.endswith(e)):
+        return CloudType.AZURE
+    elif any(e for e in DATABRICKS_AWS_DOMAINS if host.endswith(e)):
+        return CloudType.AWS
+    else:
+        return None
+
+
+def get_databricks_oidc_url(hostname: str):
+    maybe_scheme = "https://" if not hostname.startswith("https://") else ""
+    maybe_trailing_slash = "/" if not hostname.endswith("/") else ""
+    return f"{maybe_scheme}{hostname}{maybe_trailing_slash}{OIDC_REDIRECTOR_PATH}"
+
+
+class OAuthEndpointCollection(ABC):
+    @abstractmethod
+    def get_scopes_mapping(self, scopes: List[str]) -> List[str]:
+        raise NotImplementedError()
+
+    # Endpoint for oauth2 authorization  e.g https://idp.example.com/oauth2/v2.0/authorize
+    @abstractmethod
+    def get_authorization_url(self, hostname: str) -> str:
+        raise NotImplementedError()
+
+    # Endpoint for well-known openid configuration e.g https://idp.example.com/oauth2/.well-known/openid-configuration
+    @abstractmethod
+    def get_openid_config_url(self, hostname: str) -> str:
+        raise NotImplementedError()
+
+
+class AzureOAuthEndpointCollection(OAuthEndpointCollection):
+    DATATRICKS_AZURE_APP = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+
+    def get_scopes_mapping(self, scopes: List[str]) -> List[str]:
+        # There is no corresponding scopes in Azure, instead, access control will be delegated to Databricks
+        tenant_id = os.getenv(
+            "DATABRICKS_AZURE_TENANT_ID",
+            AzureOAuthEndpointCollection.DATATRICKS_AZURE_APP,
+        )
+        azure_scope = f"{tenant_id}/user_impersonation"
+        mapped_scopes = [azure_scope]
+        if OAuthScope.OFFLINE_ACCESS in scopes:
+            mapped_scopes.append(OAuthScope.OFFLINE_ACCESS)
+        return mapped_scopes
+
+    def get_authorization_url(self, hostname: str):
+        # We need get account specific url, which can be redirected by databricks unified oidc endpoint
+        return f"{get_databricks_oidc_url(hostname)}/oauth2/v2.0/authorize"
+
+    def get_openid_config_url(self, hostname: str):
+        return "https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration"
+
+
+class AwsOAuthEndpointCollection(OAuthEndpointCollection):
+    def get_scopes_mapping(self, scopes: List[str]) -> List[str]:
+        # No scope mapping in AWS
+        return scopes.copy()
+
+    def get_authorization_url(self, hostname: str):
+        idp_url = get_databricks_oidc_url(hostname)
+        return f"{idp_url}/oauth2/v2.0/authorize"
+
+    def get_openid_config_url(self, hostname: str):
+        idp_url = get_databricks_oidc_url(hostname)
+        return f"{idp_url}/.well-known/oauth-authorization-server"
+
+
+def get_oauth_endpoints(cloud: CloudType) -> Optional[OAuthEndpointCollection]:
+    if cloud == CloudType.AWS:
+        return AwsOAuthEndpointCollection()
+    elif cloud == CloudType.AZURE:
+        return AzureOAuthEndpointCollection()
+    else:
+        return None

src/databricks/sql/auth/oauth.py

@@ -14,17 +14,22 @@
 from requests.exceptions import RequestException
 
 from databricks.sql.auth.oauth_http_handler import OAuthHttpSingleRequestHandler
+from databricks.sql.auth.endpoint import OAuthEndpointCollection
 
 logger = logging.getLogger(__name__)
 
 
 class OAuthManager:
-    OIDC_REDIRECTOR_PATH = "oidc"
-
-    def __init__(self, port_range: List[int], client_id: str):
+    def __init__(
+        self,
+        port_range: List[int],
+        client_id: str,
+        idp_endpoint: OAuthEndpointCollection,
+    ):
         self.port_range = port_range
         self.client_id = client_id
         self.redirect_port = None
+        self.idp_endpoint = idp_endpoint
 
     @staticmethod
     def __token_urlsafe(nbytes=32):
@@ -34,14 +39,14 @@ def __token_urlsafe(nbytes=32):
     def __get_redirect_url(redirect_port: int):
         return f"http://localhost:{redirect_port}"
 
-    @staticmethod
-    def __fetch_well_known_config(idp_url: str):
-        known_config_url = f"{idp_url}/.well-known/oauth-authorization-server"
+    def __fetch_well_known_config(self, hostname: str):
+        known_config_url = self.idp_endpoint.get_openid_config_url(hostname)
+
         try:
             response = requests.get(url=known_config_url)
         except RequestException as e:
             logger.error(
-                f"Unable to fetch OAuth configuration from {idp_url}.\n"
+                f"Unable to fetch OAuth configuration from {known_config_url}.\n"
                 "Verify it is a valid workspace URL and that OAuth is "
                 "enabled on this account."
             )
@@ -50,7 +55,7 @@ def __fetch_well_known_config(idp_url: str):
         if response.status_code != 200:
             msg = (
                 f"Received status {response.status_code} OAuth configuration from "
-                f"{idp_url}.\n Verify it is a valid workspace URL and "
+                f"{known_config_url}.\n Verify it is a valid workspace URL and "
                 "that OAuth is enabled on this account."
             )
             logger.error(msg)
@@ -59,18 +64,12 @@ def __fetch_well_known_config(idp_url: str):
             return response.json()
         except requests.exceptions.JSONDecodeError as e:
             logger.error(
-                f"Unable to decode OAuth configuration from {idp_url}.\n"
+                f"Unable to decode OAuth configuration from {known_config_url}.\n"
                 "Verify it is a valid workspace URL and that OAuth is "
                 "enabled on this account."
             )
             raise e
 
-    @staticmethod
-    def __get_idp_url(host: str):
-        maybe_scheme = "https://" if not host.startswith("https://") else ""
-        maybe_trailing_slash = "/" if not host.endswith("/") else ""
-        return f"{maybe_scheme}{host}{maybe_trailing_slash}{OAuthManager.OIDC_REDIRECTOR_PATH}"
-
     @staticmethod
     def __get_challenge():
         verifier_string = OAuthManager.__token_urlsafe(32)
@@ -154,8 +153,7 @@ def __send_token_request(token_request_url, data):
         return response.json()
 
     def __send_refresh_token_request(self, hostname, refresh_token):
-        idp_url = OAuthManager.__get_idp_url(hostname)
-        oauth_config = OAuthManager.__fetch_well_known_config(idp_url)
+        oauth_config = self.__fetch_well_known_config(hostname)
         token_request_url = oauth_config["token_endpoint"]
         client = oauthlib.oauth2.WebApplicationClient(self.client_id)
         token_request_body = client.prepare_refresh_body(
@@ -215,14 +213,15 @@ def check_and_refresh_access_token(
         return fresh_access_token, fresh_refresh_token, True
 
     def get_tokens(self, hostname: str, scope=None):
-        idp_url = self.__get_idp_url(hostname)
-        oauth_config = self.__fetch_well_known_config(idp_url)
+        oauth_config = self.__fetch_well_known_config(hostname)
         # We are going to override oauth_config["authorization_endpoint"] use the
         # /oidc redirector on the hostname, which may inject additional parameters.
-        auth_url = f"{hostname}oidc/v1/authorize"
+        auth_url = self.idp_endpoint.get_authorization_url(hostname)
+
         state = OAuthManager.__token_urlsafe(16)
         (verifier, challenge) = OAuthManager.__get_challenge()
         client = oauthlib.oauth2.WebApplicationClient(self.client_id)
+
         try:
             auth_response = self.__get_authorization_code(
                 client, auth_url, scope, state, challenge

src/databricks/sql/experimental/oauth_persistence.py

@@ -27,6 +27,17 @@ def read(self, hostname: str) -> Optional[OAuthToken]:
         pass
 
 
+class OAuthPersistenceCache(OAuthPersistence):
+    def __init__(self):
+        self.tokens = {}
+
+    def persist(self, hostname: str, oauth_token: OAuthToken):
+        self.tokens[hostname] = oauth_token
+
+    def read(self, hostname: str) -> Optional[OAuthToken]:
+        return self.tokens.get(hostname)
+
+
 # Note this is only intended to be used for development
 class DevOnlyFilePersistence(OAuthPersistence):
     def __init__(self, file_path):

tests/unit/test_auth.py

@@ -1,8 +1,15 @@
 import unittest
+import pytest
+from typing import Optional
+from unittest.mock import patch
 
 from databricks.sql.auth.auth import AccessTokenAuthProvider, BasicAuthProvider, AuthProvider, ExternalAuthProvider
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
+from databricks.sql.auth.oauth import OAuthManager
+from databricks.sql.auth.authenticators import DatabricksOAuthProvider
+from databricks.sql.auth.endpoint import CloudType, AwsOAuthEndpointCollection, AzureOAuthEndpointCollection
 from databricks.sql.auth.authenticators import CredentialsProvider, HeaderFactory
+from databricks.sql.experimental.oauth_persistence import OAuthPersistenceCache
 
 
 class Auth(unittest.TestCase):
@@ -38,6 +45,39 @@ def test_noop_auth_provider(self):
         self.assertEqual(len(http_request.keys()), 1)
         self.assertEqual(http_request['myKey'], 'myVal')
 
+    @patch.object(OAuthManager, "check_and_refresh_access_token")
+    @patch.object(OAuthManager, "get_tokens")
+    def test_oauth_auth_provider(self, mock_get_tokens, mock_check_and_refresh):
+        client_id = "mock-id"
+        scopes = ["offline_access", "sql"]
+        access_token = "mock_token"
+        refresh_token = "mock_refresh_token"
+        mock_get_tokens.return_value = (access_token, refresh_token)
+        mock_check_and_refresh.return_value = (access_token, refresh_token, False)
+
+        params = [(CloudType.AWS, "foo.cloud.databricks.com", AwsOAuthEndpointCollection, "offline_access sql"),
+                  (CloudType.AZURE, "foo.1.azuredatabricks.net", AzureOAuthEndpointCollection,
+                   f"{AzureOAuthEndpointCollection.DATATRICKS_AZURE_APP}/user_impersonation offline_access")]
+
+        for cloud_type, host, expected_endpoint_type, expected_scopes in params:
+            with self.subTest(cloud_type.value):
+                oauth_persistence = OAuthPersistenceCache()
+                auth_provider = DatabricksOAuthProvider(hostname=host,
+                                                        oauth_persistence=oauth_persistence,
+                                                        redirect_port_range=[8020],
+                                                        client_id=client_id,
+                                                        scopes=scopes)
+
+                self.assertIsInstance(auth_provider.oauth_manager.idp_endpoint, expected_endpoint_type)
+                self.assertEqual(auth_provider.oauth_manager.port_range, [8020])
+                self.assertEqual(auth_provider.oauth_manager.client_id, client_id)
+                self.assertEqual(oauth_persistence.read(host).refresh_token, refresh_token)
+                mock_get_tokens.assert_called_with(hostname=host, scope=expected_scopes)
+
+                headers = {}
+                auth_provider.add_headers(headers)
+                self.assertEqual(headers['Authorization'], f"Bearer {access_token}")
+
     def test_external_provider(self):
         class MyProvider(CredentialsProvider):
                 def auth_type(self) -> str: