Skip to content

Conversation

madhav-db
Copy link
Contributor

@madhav-db madhav-db commented Sep 5, 2025

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

Copy link

github-actions bot commented Sep 5, 2025

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).

@madhav-db madhav-db changed the title Token federation for python driver Add Token Federation Support for Databricks SQL Python Driver Sep 8, 2025
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.

1 participant