Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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.
Expand Down
21 changes: 20 additions & 1 deletion src/databricks/sql/auth/token_federation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
66 changes: 63 additions & 3 deletions tests/unit/test_token_federation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
[
Expand Down