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

Bắt đầu tại src/databricks/sql/auth/token_federation.py, ở _exchange_token quanh các dòng 191-194, sau đó lần theo handler quanh dòng 148 và entry point bao quanh trong src/databricks/sql/auth/auth.py:68. Tái hiện một exchange trả về OAuth error body và xác minh rằng warning hiển thị lý do lỗi của endpoint mà không ghi log các token, trong khi fallback sang external token vẫn không thay đổi.

Summary

When token exchange fails, _exchange_token raises KeyError: 'access_token' instead of surfacing the error the endpoint actually returned. The caller catches it and logs the KeyError, so the log line reports the name of a missing dict key rather than why the exchange was rejected:

Token exchange failed, using external token: 'access_token'

There is no way to tell from that whether the exchange was misconfigured, unauthorized, or unsupported.

Versions

Reproduced on databricks-sql-connector==4.3.0. The relevant code is byte-identical in v4.4.0 and on main today, so this is not fixed in a later release.

Where

src/databricks/sql/auth/token_federation.py on main:

191:        token_response = json.loads(response.data.decode())
192:
193:        return Token(
194:            token_response["access_token"], token_response.get("token_type", "Bearer")
195:        )

token_type is read defensively with .get(); access_token is not. When the endpoint returns an OAuth error body ({"error": ..., "error_description": ...}) rather than a token, line 194 raises KeyError.

That propagates to the handler at line 148:

147:            except Exception as e:
148:                logger.warning("Token exchange failed, using external token: %s", e)

str(KeyError("access_token")) renders as 'access_token', which is what reaches the log. The error and error_description from the response body are never read and are lost.

Reproduction

1. Connect using an OAuth access token issued by an identity provider whose iss host differs from the workspace host — for example an Microsoft Entra ID token for the Azure Databricks resource (2ff814a6-3304-4ab8-85cb-cd0e6f879c1d).
2. TokenFederationProvider wraps every provider unconditionally (src/databricks/sql/auth/auth.py:68, "Always wrap with token federation"), and _should_exchange_token returns True because the issuer host does not match the workspace host, so an exchange is always attempted.
3. Against a workspace where that exchange is not accepted, every connection logs the message above.

Functionally this is harmless — the fallback to the external token works correctly and queries succeed. The problem is purely diagnostic: the warning fires on every connection and gives no actionable information.

Suggested fix

Read the response defensively and raise something that names the actual failure, without logging the token itself:

token_response = json.loads(response.data.decode())

if "access_token" not in token_response:
    error = token_response.get("error", "unknown_error")
    description = token_response.get("error_description", "")
    raise RuntimeError(f"Token exchange rejected by {token_url}: {error} {description}".strip())

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

The existing except Exception at line 147 would then log the endpoint's own reason, and the graceful fallback behaviour is unchanged.

A non-JSON or non-2xx response would also currently surface as a confusing JSONDecodeError; checking the status code before parsing would cover that case too.