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

Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>

Author: peco-engineer-bot[bot]

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

@@ -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