fix: surface the token-exchange error instead of KeyError('access_token')

When token federation exchange fails, `_exchange_token` read
`token_response["access_token"]` unconditionally. An OAuth error body
(`{"error": ..., "error_description": ...}`) therefore raised
`KeyError('access_token')`, and the handler in `_get_token` logged the
name of the missing key:

    Token exchange failed, using external token: 'access_token'

The endpoint's own reason was never read. Since `TokenFederationProvider`
wraps every provider and `_should_exchange_token` returns True whenever
the issuer host differs from the workspace host, this warning fires on
every connection for cross-issuer tokens with no way to tell whether the
exchange was misconfigured, unauthorized, or unsupported.

Check for `access_token` before reading it and raise a ValueError that
names the endpoint, the HTTP status, and the returned `error` /
`error_description`. A non-JSON body now reports the endpoint and status
rather than surfacing a JSONDecodeError. The response carries no token in
either case, so nothing sensitive is exposed.

The fallback to the external token is unchanged: `_get_token` still
catches the exception and connections keep working.

Resolves #904

Signed-off-by: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com>

Author: TangoEnSkai

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Release History
 
 # Unreleased
+- Fix: a rejected token-federation exchange now reports the reason the endpoint gave. `_exchange_token` raised `KeyError: 'access_token'` on an OAuth error body, so the connection logged `Token exchange failed, using external token: 'access_token'` and the endpoint's `error` / `error_description` were discarded. It now raises a `ValueError` naming the endpoint, the HTTP status, and the returned error, and a non-JSON body reports the endpoint and status instead of surfacing a `JSONDecodeError`. The graceful fallback to the external token is unchanged ([#904](https://github.com/databricks/databricks-sql-python/issues/904))
 - Kernel backend (`use_kernel=True`): OAuth **M2M with a JWT private-key client assertion** (RFC 7523) is now supported. Pass `oauth_client_id` + `oauth_jwt_key_file` + `oauth_jwt_kid` (with optional `oauth_jwt_passphrase` for an encrypted PKCS#8 key, `oauth_jwt_algorithm` defaulting to `RS256`, `oauth_scopes`, and `token_url` for the IdP token endpoint) and the connector routes them to the kernel's `auth_type="oauth-m2m-jwt"`, which signs a short-lived assertion with the private key instead of sending a client secret. The kernel owns the token lifecycle. A private-key file is treated as unambiguous JWT M2M intent and is mutually exclusive with `oauth_client_secret` / `credentials_provider` (both raise `NotSupportedError`). Verified end-to-end against an Azure Databricks workspace with the service principal's public certificate registered on its Entra ID app registration. Requires `databricks-sql-kernel >= 0.2.0` with JWT support.
 - Kernel backend (`use_kernel=True`): OAuth U2M with `auth_type="databricks-oauth"` now forwards the connector's `databricks-sql-python` OAuth-app bundle (`client_id` + `sql offline_access` scopes + redirect port) into the kernel, so a bare U2M connection authenticates as `databricks-sql-python` — parity with the Thrift path — instead of inheriting the kernel's own `databricks-sql-connector` default. A caller-supplied `oauth_client_id` (with its coupled `oauth_redirect_port`) is honored, as is a caller-supplied `oauth_scopes`; absent one, the connector default (`sql offline_access`) is forwarded. Note: the kernel binds a single U2M redirect port, so unlike the Thrift path (which tries the full `8020..8024` range) the kernel path uses only one port and does not fall back to the next port if it is already bound — pass `oauth_redirect_port` (with `oauth_client_id`) to pick a free one on a port collision (PECOBLR-4040)
 - Kernel backend (`use_kernel=True`): **Azure Entra (Azure AD) service-principal M2M is now supported.** `auth_type="azure-sp-m2m"` forwards `azure_client_id` / `azure_client_secret`; the kernel is the Azure-aware auth core — it builds the Entra v2.0 token endpoint and the `{app_id}/.default` scope, and **auto-discovers the tenant** from the workspace's `/aad/auth` redirect when `azure_tenant_id` is omitted (matching Thrift). The `Authorization` bearer is the Databricks-audience data token, which alone authenticates a workspace-member SP. Set `azure_workspace_resource_id` and the kernel also sends the Azure SP management token (`X-Databricks-Azure-SP-Management-Token`) + `X-Databricks-Azure-Workspace-Resource-Id` header (matching the JDBC driver), so a service principal with an Azure RBAC role but no workspace membership can authenticate; omit it and no ARM management-scope token is fetched. Azure AD **U2M** (`auth_type="azure-oauth"`) now routes to the kernel's OAuth U2M flow, identically to `auth_type="databricks-oauth"`: the kernel runs the in-house workspace-federated browser flow, which Azure workspaces support (the workspace federates login to Entra). It forwards the connector's `databricks-sql-python` OAuth app, not the Thrift Azure app (`96eecda7` / port 8030), which is registered for Thrift's direct-Entra flow the kernel does not perform (PECOBLR-4141; PECOBLR-4120)

src/databricks/sql/auth/token_federation.py

@@ -188,7 +188,26 @@ def _exchange_token(self, access_token: str) -> Token:
             HttpMethod.POST, url=token_url, body=body, headers=headers
         )
 
-        token_response = json.loads(response.data.decode())
+        status = getattr(response, "status", None)
+
+        try:
+            token_response = json.loads(response.data.decode())
+        except (ValueError, UnicodeDecodeError) as e:
+            raise ValueError(
+                f"Token exchange at {token_url} returned a non-JSON response "
+                f"(HTTP {status})"
+            ) from e
+
+        if "access_token" not in token_response:
+            # An OAuth error body carries the reason the exchange was refused.
+            # Surface it instead of letting a KeyError hide it. The response
+            # holds no token in this case, so nothing sensitive is exposed.
+            error = token_response.get("error", "unknown_error")
+            description = token_response.get("error_description", "")
+            raise ValueError(
+                f"Token exchange at {token_url} was rejected (HTTP {status}): "
+                f"{error} {description}".strip()
+            )
 
         return Token(
             token_response["access_token"], token_response.get("token_type", "Bearer")

tests/unit/test_token_federation.py

@@ -194,15 +194,75 @@ def test_exchange_token_success(self, token_federation_provider, mock_http_clien
         assert parsed_body["client_id"][0] == "test-client-id"
 
     def test_exchange_token_failure(self, token_federation_provider, mock_http_client):
-        """Test token exchange failure handling."""
+        """An OAuth error body is surfaced instead of a bare KeyError."""
         mock_response = Mock()
-        mock_response.data = b'{"error": "invalid_request"}'
+        mock_response.data = (
+            b'{"error": "invalid_request", '
+            b'"error_description": "subject token is not supported"}'
+        )
         mock_response.status = 400
         mock_http_client.request.return_value = mock_response
 
-        with pytest.raises(KeyError):  # Will raise KeyError due to missing access_token
+        with pytest.raises(ValueError) as exc_info:
+            token_federation_provider._exchange_token("external-token-123")
+
+        message = str(exc_info.value)
+        assert "invalid_request" in message
+        assert "subject token is not supported" in message
+        assert "400" in message
+        assert "https://test.databricks.com/oidc/v1/token" in message
+
+    def test_exchange_token_failure_without_error_description(
+        self, token_federation_provider, mock_http_client
+    ):
+        """A body carrying no error fields still names the failing endpoint."""
+        mock_response = Mock()
+        mock_response.data = b"{}"
+        mock_response.status = 401
+        mock_http_client.request.return_value = mock_response
+
+        with pytest.raises(ValueError) as exc_info:
+            token_federation_provider._exchange_token("external-token-123")
+
+        assert "unknown_error" in str(exc_info.value)
+
+    def test_exchange_token_non_json_response(
+        self, token_federation_provider, mock_http_client
+    ):
+        """A non-JSON body reports the endpoint rather than a JSONDecodeError."""
+        mock_response = Mock()
+        mock_response.data = b"<html>502 Bad Gateway</html>"
+        mock_response.status = 502
+        mock_http_client.request.return_value = mock_response
+
+        with pytest.raises(ValueError) as exc_info:
             token_federation_provider._exchange_token("external-token-123")
 
+        message = str(exc_info.value)
+        assert "non-JSON" in message
+        assert "502" in message
+
+    def test_exchange_token_failure_keeps_external_token_fallback(
+        self, token_federation_provider, mock_http_client, mock_external_provider
+    ):
+        """A rejected exchange still falls back to the external token."""
+        external_token = create_jwt_token(
+            issuer="https://login.microsoftonline.com/tenant-id/"
+        )
+        mock_external_provider.add_headers.side_effect = (
+            lambda headers: headers.update({"Authorization": f"Bearer {external_token}"})
+        )
+
+        mock_response = Mock()
+        mock_response.data = b'{"error": "invalid_request"}'
+        mock_response.status = 400
+        mock_http_client.request.return_value = mock_response
+
+        request_headers: dict = {}
+        token_federation_provider.add_headers(request_headers)
+
+        assert request_headers["Authorization"] == f"Bearer {external_token}"
+
     @pytest.mark.parametrize(
         "external_issuer,should_exchange",
         [