diff --git a/src/databricks/sql/auth/auth.py b/src/databricks/sql/auth/auth.py index a4d4d6f2e..fadf2e63d 100755 --- a/src/databricks/sql/auth/auth.py +++ b/src/databricks/sql/auth/auth.py @@ -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): @@ -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 ) @@ -112,23 +116,51 @@ 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") + + if caller_client_id: + client_id = caller_client_id + # Caller owns the bundle: forward their scopes verbatim. + scopes = oauth_scopes or PYSQL_OAUTH_SCOPES + # 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 + # The caller relies on the driver's default client_id, whose OAuth app + # only has the app-specific default scopes/range registered. A caller- + # supplied oauth_scopes must NOT override the driver's app-specific + # default scopes here, just as a caller-supplied oauth_redirect_port must + # NOT override the default range (doing so risks a redirect_uri_mismatch + # or scope error, PECOBLR-4039); both only take effect alongside a + # caller-supplied oauth_client_id. + scopes = PYSQL_OAUTH_SCOPES + redirect_port_range = default_redirect_port_range + 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"), diff --git a/tests/unit/test_auth.py b/tests/unit/test_auth.py index d1b941208..449fb44dd 100644 --- a/tests/unit/test_auth.py +++ b/tests/unit/test_auth.py @@ -12,6 +12,8 @@ from databricks.sql.auth.auth import ( get_python_sql_connector_auth_provider, PYSQL_OAUTH_CLIENT_ID, + PYSQL_OAUTH_REDIRECT_PORT_RANGE, + PYSQL_OAUTH_SCOPES, ) from databricks.sql.auth.oauth import OAuthManager, Token, ClientCredentialsTokenSource from databricks.sql.auth.authenticators import ( @@ -207,6 +209,96 @@ 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]) + + @patch.object(DatabricksOAuthProvider, "_initial_get_token") + def test_get_python_sql_connector_u2m_default_client_id_ignores_redirect_port( + self, mock__initial_get_token + ): + # AUTH-013 / PECOBLR-4039: a caller relying on the driver's DEFAULT + # client_id must stay pinned to the driver's registered redirect port + # range even if oauth_redirect_port is supplied - the driver's OAuth app + # only registers 8020-8024, so honoring an arbitrary port would produce a + # redirect_uri_mismatch. oauth_redirect_port only applies with a + # caller-supplied oauth_client_id. + hostname = "foo.cloud.databricks.com" + kwargs = {"oauth_redirect_port": 9000} + 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, PYSQL_OAUTH_CLIENT_ID) + self.assertEqual( + provider.oauth_manager.port_range, PYSQL_OAUTH_REDIRECT_PORT_RANGE + ) + + @patch.object(DatabricksOAuthProvider, "_initial_get_token") + def test_get_python_sql_connector_u2m_default_client_id_ignores_scopes( + self, mock__initial_get_token + ): + # AUTH-013 / PECOBLR-4039: a caller relying on the driver's DEFAULT + # client_id must stay pinned to the driver's app-specific default scopes + # even if oauth_scopes is supplied - the driver's OAuth app only has its + # own default scopes registered, so honoring arbitrary scopes risks a + # scope error. oauth_scopes only applies with a caller-supplied + # oauth_client_id. + hostname = "foo.cloud.databricks.com" + kwargs = {"oauth_scopes": ["all-apis"]} + 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, PYSQL_OAUTH_CLIENT_ID) + self.assertEqual(provider._scopes_as_str, " ".join(PYSQL_OAUTH_SCOPES)) + class TestClientCredentialsTokenSource: @pytest.fixture