Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
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
41 changes: 33 additions & 8 deletions src/databricks/sql/auth/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,10 @@ def get_auth_provider(cfg: ClientContext, http_client):
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]
# Base (app-neutral) redirect port used when a caller supplies their OWN
# oauth_client_id but no redirect port: the driver must NOT pin its own
# app-specific default range in that case (see AUTH-013 / PECOBLR-4039).
PYSQL_OAUTH_BASE_REDIRECT_PORT_RANGE = [8030]


def normalize_host_name(hostname: str):
Expand All @@ -102,7 +106,7 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **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(
default_client_id, default_redirect_port_range = get_client_id_and_redirect_port(
auth_type == AuthType.AZURE_OAUTH.value
)

Expand All @@ -112,23 +116,44 @@ def get_python_sql_connector_auth_provider(hostname: str, http_client, **kwargs)
"Please use OAuth or access token instead."
)

# A caller who supplies their OWN oauth_client_id owns the rest of the U2M
# bundle: the driver forwards their scopes and redirect port verbatim and must
# NOT substitute its own app-specific defaults (see AUTH-013 / PECOBLR-4039).
# Only when the caller relies on the driver's default client_id do the
# driver's app-specific default scopes/port range apply.
caller_client_id = kwargs.get("oauth_client_id")
oauth_redirect_port = kwargs.get("oauth_redirect_port")
oauth_scopes = kwargs.get("oauth_scopes")

scopes = oauth_scopes or PYSQL_OAUTH_SCOPES
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
if caller_client_id:
client_id = caller_client_id
# Foreign client_id with no explicit port falls through to the base
# (app-neutral) default, NOT the driver's app-specific range.
redirect_port_range = (
[oauth_redirect_port]
if oauth_redirect_port
else PYSQL_OAUTH_BASE_REDIRECT_PORT_RANGE
)
else:
client_id = default_client_id
redirect_port_range = (
[oauth_redirect_port] if oauth_redirect_port else default_redirect_port_range
Comment thread
peco-review-bot[bot] marked this conversation as resolved.
Outdated
)

cfg = ClientContext(
hostname=normalize_host_name(hostname),
auth_type=auth_type,
access_token=kwargs.get("access_token"),
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 client_id,
oauth_scopes=scopes,
oauth_client_id=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
),
oauth_redirect_port_range=redirect_port_range,
oauth_persistence=kwargs.get("experimental_oauth_persistence"),
credentials_provider=kwargs.get("credentials_provider"),
identity_federation_client_id=kwargs.get("identity_federation_client_id"),
Expand Down
44 changes: 44 additions & 0 deletions tests/unit/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,50 @@ def test_get_python_sql_connector_default_auth(self, mock__initial_get_token):

self.assertEqual(auth_provider.external_provider._client_id, PYSQL_OAUTH_CLIENT_ID)

@patch.object(DatabricksOAuthProvider, "_initial_get_token")
def test_get_python_sql_connector_u2m_explicit_bundle_override(
self, mock__initial_get_token
):
# AUTH-013 Case 1: a caller who supplies their OWN client_id owns the rest
# of the U2M bundle - the caller's client_id, scopes and redirect port are
# forwarded verbatim, with NO driver default substitution.
hostname = "foo.cloud.databricks.com"
kwargs = {
"oauth_client_id": "test-custom-u2m-app",
"oauth_scopes": ["all-apis"],
"oauth_redirect_port": 8099,
}
mock_http_client = MagicMock()
auth_provider = get_python_sql_connector_auth_provider(
hostname, mock_http_client, **kwargs
)

provider = auth_provider.external_provider
self.assertEqual(type(provider).__name__, "DatabricksOAuthProvider")
self.assertEqual(provider._client_id, "test-custom-u2m-app")
self.assertEqual(provider.oauth_manager.port_range, [8099])
# Caller's scope set forwarded verbatim, NOT the driver default "sql offline_access".
self.assertEqual(provider._scopes_as_str, "all-apis")

@patch.object(DatabricksOAuthProvider, "_initial_get_token")
def test_get_python_sql_connector_u2m_foreign_client_id_no_port(
self, mock__initial_get_token
):
# AUTH-013 Case 2: a foreign client_id with NO redirect port must NOT be
# pinned to the driver's own app-specific default port range (8020-8024).
# It falls through to the base kernel default [8030].
hostname = "foo.cloud.databricks.com"
kwargs = {"oauth_client_id": "test-custom-u2m-app"}
mock_http_client = MagicMock()
auth_provider = get_python_sql_connector_auth_provider(
hostname, mock_http_client, **kwargs
)

provider = auth_provider.external_provider
self.assertEqual(type(provider).__name__, "DatabricksOAuthProvider")
self.assertEqual(provider._client_id, "test-custom-u2m-app")
self.assertEqual(provider.oauth_manager.port_range, [8030])


class TestClientCredentialsTokenSource:
@pytest.fixture
Expand Down
Loading