Fix token federation warnings by making it opt-in and improving error handling [Issue #702] #710
Conversation
madhav-db
had a problem deploying
to
azure-prod
GitHub Actions
Failure
madhav-db
had a problem deploying
to
azure-prod
GitHub Actions
Failure
|
Thanks for your contribution! To satisfy the DCO policy in our contributing guide every commit message must include a sign-off message. One or more of your commits is missing this message. You can reword previous commit messages with an interactive rebase ( |
There was a problem hiding this comment.
Pull Request Overview
This PR fixes token federation warnings by making it opt-in and improving error handling. Previously, token federation was always enabled for all authentication methods, causing warnings for users who weren't using identity federation. Now, token federation is only enabled when identity_federation_client_id is explicitly provided.
Key Changes:
- Token federation is now opt-in via the
identity_federation_client_idparameter - Added custom exception classes for better error handling and categorization
- Enhanced logging with appropriate levels (debug/info/warning) based on error type
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| src/databricks/sql/auth/auth.py | Modified to only wrap providers with TokenFederationProvider when identity_federation_client_id is provided |
| src/databricks/sql/auth/token_federation.py | Added custom exception classes and improved error handling with specific exceptions for different HTTP status codes |
| tests/unit/test_auth.py | Updated tests to verify token federation is only enabled with identity_federation_client_id and reformatted some lines |
| tests/unit/test_token_federation.py | Added tests for new exception types and fallback behavior on 404/401/403 errors |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if response.status == 404: | ||
| raise TokenExchangeNotAvailableError( | ||
| "Token exchange endpoint not found. Token federation may not be enabled for this workspace." | ||
| ) |
There was a problem hiding this comment.
[nitpick] The error messages are hardcoded strings. Consider defining them as class constants (e.g., ERROR_MSG_404, ERROR_MSG_401_403) for better maintainability and consistency, especially since these messages are validated in tests.
madhav-db
changed the title
|
Fix for issue #702 |
| # Always wrap with token federation (falls back gracefully if not needed) | ||
| if base_provider: | ||
| # Wrap with token federation only if explicitly enabled via identity_federation_client_id | ||
| if base_provider and cfg.identity_federation_client_id: |
There was a problem hiding this comment.
For account wide token federation, this client id is not present? IIUC, this will completely skip the token federation
| logger.warning( | ||
| "Token exchange failed due to authentication error. Using external token directly. " | ||
| "Error: %s. If this persists, verify your identity_federation_client_id configuration.", | ||
| e, |
| class TokenFederationError(Exception): | ||
| """Base exception for token federation errors.""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class TokenExchangeNotAvailableError(TokenFederationError): | ||
| """Raised when token exchange endpoint is not available (404).""" | ||
|
|
||
| pass | ||
|
|
||
|
|
||
| class TokenExchangeAuthenticationError(TokenFederationError): | ||
| """Raised when token exchange fails due to authentication issues (401/403).""" | ||
|
|
||
| pass | ||
|
|
||
|
|
There was a problem hiding this comment.
Why are separate error classes being created for standard Http errors? Like 404, 401, 403 etc are just Http errors, so don't think need to create something different for them, just catch and log the normal errors
| HttpMethod.POST, url=token_url, body=body, headers=headers | ||
| ) | ||
|
|
||
| # Check response status code |
There was a problem hiding this comment.
Can you use the response with the with context, to handler auto close.
|
|
||
| # Check response status code | ||
| if response.status == 404: | ||
| raise TokenExchangeNotAvailableError( |
There was a problem hiding this comment.
No need to manually handle everything just do response.raise_for_status, it auto handles all these basic errors
| error_detail = ( | ||
| response.data.decode() if response.data else "No error details" | ||
| ) | ||
| raise TokenExchangeAuthenticationError( |
There was a problem hiding this comment.
Same comment, basic http errors don't need to be handled, just do response.raise_for_status
| token issuer with the Databricks workspace hostname. If exchange fails, it gracefully | ||
| falls back to using the external token directly. | ||
| Note: Token federation must be explicitly enabled by providing the |
There was a problem hiding this comment.
How does it work with account wide token federation then?
What type of PR is this?
Fixes token federation warnings that were appearing for users not using identity federation. Token federation is now opt-in and only enabled when
identity_federation_client_idis explicitly provided.Description
How is this tested?
Related Tickets & Documents