[coverage] Conformance findings: AUTH-013 (#909)

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"),

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