Skip to content

Commit 3265132

Browse files
committed
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>
1 parent 13e8af4 commit 3265132

3 files changed

Lines changed: 86 additions & 4 deletions

File tree

CHANGELOG.md

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
# Release History
22

3+
# Unreleased
4+
- 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))
5+
36
# 4.5.0 (2026-09-01)
47
- Upgrade Databricks SQL Kernel to 1.0.0.
58
- Add JWT private-key M2M and Azure Entra authentication for kernel connections.

src/databricks/sql/auth/token_federation.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,26 @@ def _exchange_token(self, access_token: str) -> Token:
188188
HttpMethod.POST, url=token_url, body=body, headers=headers
189189
)
190190

191-
token_response = json.loads(response.data.decode())
191+
status = getattr(response, "status", None)
192+
193+
try:
194+
token_response = json.loads(response.data.decode())
195+
except (ValueError, UnicodeDecodeError) as e:
196+
raise ValueError(
197+
f"Token exchange at {token_url} returned a non-JSON response "
198+
f"(HTTP {status})"
199+
) from e
200+
201+
if "access_token" not in token_response:
202+
# An OAuth error body carries the reason the exchange was refused.
203+
# Surface it instead of letting a KeyError hide it. The response
204+
# holds no token in this case, so nothing sensitive is exposed.
205+
error = token_response.get("error", "unknown_error")
206+
description = token_response.get("error_description", "")
207+
raise ValueError(
208+
f"Token exchange at {token_url} was rejected (HTTP {status}): "
209+
f"{error} {description}".strip()
210+
)
192211

193212
return Token(
194213
token_response["access_token"], token_response.get("token_type", "Bearer")

tests/unit/test_token_federation.py

Lines changed: 63 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -194,15 +194,75 @@ def test_exchange_token_success(self, token_federation_provider, mock_http_clien
194194
assert parsed_body["client_id"][0] == "test-client-id"
195195

196196
def test_exchange_token_failure(self, token_federation_provider, mock_http_client):
197-
"""Test token exchange failure handling."""
197+
"""An OAuth error body is surfaced instead of a bare KeyError."""
198198
mock_response = Mock()
199-
mock_response.data = b'{"error": "invalid_request"}'
199+
mock_response.data = (
200+
b'{"error": "invalid_request", '
201+
b'"error_description": "subject token is not supported"}'
202+
)
200203
mock_response.status = 400
201204
mock_http_client.request.return_value = mock_response
202205

203-
with pytest.raises(KeyError): # Will raise KeyError due to missing access_token
206+
with pytest.raises(ValueError) as exc_info:
207+
token_federation_provider._exchange_token("external-token-123")
208+
209+
message = str(exc_info.value)
210+
assert "invalid_request" in message
211+
assert "subject token is not supported" in message
212+
assert "400" in message
213+
assert "https://test.databricks.com/oidc/v1/token" in message
214+
215+
def test_exchange_token_failure_without_error_description(
216+
self, token_federation_provider, mock_http_client
217+
):
218+
"""A body carrying no error fields still names the failing endpoint."""
219+
mock_response = Mock()
220+
mock_response.data = b"{}"
221+
mock_response.status = 401
222+
mock_http_client.request.return_value = mock_response
223+
224+
with pytest.raises(ValueError) as exc_info:
225+
token_federation_provider._exchange_token("external-token-123")
226+
227+
assert "unknown_error" in str(exc_info.value)
228+
229+
def test_exchange_token_non_json_response(
230+
self, token_federation_provider, mock_http_client
231+
):
232+
"""A non-JSON body reports the endpoint rather than a JSONDecodeError."""
233+
mock_response = Mock()
234+
mock_response.data = b"<html>502 Bad Gateway</html>"
235+
mock_response.status = 502
236+
mock_http_client.request.return_value = mock_response
237+
238+
with pytest.raises(ValueError) as exc_info:
204239
token_federation_provider._exchange_token("external-token-123")
205240

241+
message = str(exc_info.value)
242+
assert "non-JSON" in message
243+
assert "502" in message
244+
245+
def test_exchange_token_failure_keeps_external_token_fallback(
246+
self, token_federation_provider, mock_http_client, mock_external_provider
247+
):
248+
"""A rejected exchange still falls back to the external token."""
249+
external_token = create_jwt_token(
250+
issuer="https://login.microsoftonline.com/tenant-id/"
251+
)
252+
mock_external_provider.add_headers.side_effect = (
253+
lambda headers: headers.update({"Authorization": f"Bearer {external_token}"})
254+
)
255+
256+
mock_response = Mock()
257+
mock_response.data = b'{"error": "invalid_request"}'
258+
mock_response.status = 400
259+
mock_http_client.request.return_value = mock_response
260+
261+
request_headers: dict = {}
262+
token_federation_provider.add_headers(request_headers)
263+
264+
assert request_headers["Authorization"] == f"Bearer {external_token}"
265+
206266
@pytest.mark.parametrize(
207267
"external_issuer,should_exchange",
208268
[

0 commit comments

Comments
 (0)