Skip to content

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

Open
TangoEnSkai wants to merge 1 commit into
databricks:mainfrom
TangoEnSkai:fix/token-exchange-error-message
Open

fix: surface the token-exchange error instead of KeyError('access_token')#928
TangoEnSkai wants to merge 1 commit into
databricks:mainfrom
TangoEnSkai:fix/token-exchange-error-message

Conversation

@TangoEnSkai

Copy link
Copy Markdown

Context

TokenFederationProvider wraps every credentials provider unconditionally
(auth.py:68, "Always wrap with token federation"), and _should_exchange_token
returns True whenever the token's issuer host differs from the workspace host —
for example a Microsoft Entra ID token against an Azure Databricks workspace. An
exchange is therefore attempted on every such connection.

When the exchange endpoint refuses, it replies with an OAuth error body
({"error": ..., "error_description": ...}) rather than a token.
_exchange_token read token_response["access_token"] unconditionally, so that
reply raised KeyError('access_token'). The handler in _get_token logs the
exception, and str(KeyError("access_token")) renders as the key name:

Token exchange failed, using external token: 'access_token'

The endpoint's own error and error_description were never read, so there was
no way to tell whether the exchange was misconfigured, unauthorized, or simply
unsupported for that workspace. Note token_type on the very next line is
already read defensively with .get().

What

  • _exchange_token checks for access_token before reading it. If it is
    missing, it raises a ValueError naming the endpoint, the HTTP status, and the
    returned error / error_description.
  • A non-JSON or truncated body now raises a ValueError naming the endpoint and
    status instead of surfacing a bare JSONDecodeError.
  • tests/unit/test_token_federation.py: test_exchange_token_failure previously
    asserted pytest.raises(KeyError) — it pinned the behaviour this issue is
    about — and is updated to assert the endpoint's reason is surfaced. Three tests
    are added for the no-error-fields body, the non-JSON body, and the unchanged
    fallback to the external token.

The response carries no token on the failure paths, so nothing sensitive is
placed in the message.

Why

This is diagnostic only — the fallback to the external token works and queries
succeed, which is exactly why it is worth fixing: the warning fires on every
connection and tells the operator nothing actionable. The fix is contained to the
one function and preserves the graceful degradation, since _get_token still
catches the exception and falls through to the external token.

Completion Criteria

  • _exchange_token surfaces the endpoint's error / error_description
  • Non-JSON responses report the endpoint and status, not a JSONDecodeError
  • Fallback to the external token is unchanged, covered by a new test
  • pytest tests/unit — 531 passed, 217 skipped, 185 subtests passed
  • Reverting the source change fails exactly the three new/updated tests and
    nothing else
  • black --check src clean
  • CHANGELOG entry under # Unreleased
  • Commit signed off (DCO)

close #904

…en')

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 databricks#904

Signed-off-by: TangoEnSkai <21152231+TangoEnSkai@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Token federation: _exchange_token raises KeyError('access_token') on error responses, discarding the real failure reason

1 participant