Add Token Federation Support for Databricks SQL Python Driver

[madhav-db]

What type of PR is this?

• Refactor
• Feature
• Bug Fix
• Other

Description

This PR implements token federation functionality for the databricks-sql-python driver, enabling seamless integration with external Identity Providers (IdPs) like Azure AD, Okta, and others.
Token federation allows users to authenticate with external IdPs and automatically exchange those tokens for Databricks in-house tokens when needed. This enables:

• Workflow-Wide (M2M) Token Federation: Service principal authentication with external IdPs
• Account-Wide (U2M) Token Federation: User authentication through external IdPs

Flow:

1. Token Analysis: Checks if the external token's issuer differs from the Databricks host
2. Automatic Exchange: If different, exchanges the token via /oidc/v1/token endpoint using OAuth 2.0 token exchange flow
3. Caching: Caches exchanged tokens until expiry
4. Fallback: Uses the original external token if exchange fails

How is this tested?

• Unit tests
• E2E Tests
• Manually
• N/A

Extensive testing was performed covering:

1. M2M (Machine-to-Machine) Token Federation
   - External service principal tokens from Azure AD
   - Automatic token exchange with Databricks workspace
   - Authentication as service principal in Databricks
2. U2M (User-to-Machine) Token Federation
   - Browser-based OAuth flow with automatic token handling
   - Pre-obtained user tokens from external IdPs
   - Authentication as actual users in Databricks
3. Token Lifecycle Management
   - Token caching with proper expiry handling
   - Automatic refresh when tokens expire
   - Graceful fallback when exchange fails
4. Cross-Cloud Compatibility
   - Tested with GCP Databricks workspace using Azure AD tokens
   - Tested with Azure Databricks workspace
   - Verified issuer-based exchange decision logic

Related Tickets & Documents

[github-actions]

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 (git rebase -i main).

[github-actions]

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 (git rebase -i main).

[github-actions]

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 (git rebase -i main).

[github-actions]

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 (git rebase -i main).

[github-actions]

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 (git rebase -i main).

[jprakash-db]

LGTM. Thanks for making the changes

[github-actions]

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 (git rebase -i main).

[github-actions]

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 (git rebase -i main).