diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f021ed09..4b5b37d23 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # 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)) + # 4.5.0 (2026-09-01) - Upgrade Databricks SQL Kernel to 1.0.0. - Add JWT private-key M2M and Azure Entra authentication for kernel connections. diff --git a/src/databricks/sql/auth/token_federation.py b/src/databricks/sql/auth/token_federation.py index b7cfaafa0..1a864a7e5 100644 --- a/src/databricks/sql/auth/token_federation.py +++ b/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") diff --git a/tests/unit/test_token_federation.py b/tests/unit/test_token_federation.py index 9c209e894..d4b42f5b6 100644 --- a/tests/unit/test_token_federation.py +++ b/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"502 Bad Gateway" + 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", [