Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
jackyhu-db merged 8 commits intodatabricks:mainfromjackyhu-db:PECO-626
Jun 20, 2023
Merged
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletionsCHANGELOG.md
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down
16 changes: 14 additions & 2 deletionssrc/databricks/sql/auth/auth.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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


Expand DownExpand Up@@ -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):
Expand All@@ -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"),
Expand All@@ -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") orPYSQL_OAUTH_CLIENT_ID,
oauth_client_id=kwargs.get("oauth_client_id") orclient_id,
oauth_redirect_port_range=[kwargs["oauth_redirect_port"]]
if kwargs.get("oauth_client_id") and kwargs.get("oauth_redirect_port")
elsePYSQL_OAUTH_REDIRECT_PORT_RANGE,
elseredirect_port_range,
oauth_persistence=kwargs.get("experimental_oauth_persistence"),
credentials_provider=kwargs.get("credentials_provider"),
)
Expand Down
20 changes: 18 additions & 2 deletionssrc/databricks/sql/auth/authenticators.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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.
Expand DownExpand Up@@ -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
Expand Down
112 changes: 112 additions & 0 deletionssrc/databricks/sql/auth/endpoint.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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
39 changes: 19 additions & 20 deletionssrc/databricks/sql/auth/oauth.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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):
Expand All@@ -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."
)
Expand All@@ -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)
Expand All@@ -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)
Expand DownExpand Up@@ -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(
Expand DownExpand Up@@ -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
Expand Down
11 changes: 11 additions & 0 deletionssrc/databricks/sql/experimental/oauth_persistence.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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):
Expand Down
40 changes: 40 additions & 0 deletionstests/unit/test_auth.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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):
Expand DownExpand Up@@ -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:
Expand Down
Loading
[8]ページ先頭
©2009-2025 Movatter.jp