Token federation for Go driver (2/3) #291
Conversation
Adds automatic token exchange (federation) and caching capabilities: - CachedTokenProvider: Automatic token refresh with 5min buffer - FederationProvider: Auto-detects and exchanges external JWT tokens - Supports both user federation and SP-wide (M2M) federation - Graceful fallback if token exchange unavailable - Connector functions: WithFederatedTokenProvider, WithFederatedTokenProviderAndClientID - Azure domain list updates for staging/dev environments Token exchange follows RFC 8693 standard. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
madhav-db
requested review from
deeksha-db,
gopalldb,
jackyhu-db,
jadewang-db,
jayantsing-db,
jprakash-db,
samikshya-db,
shivam2680 and
vikrantpuppala
as code owners
There was a problem hiding this comment.
Pull request overview
This PR implements token federation for the Go driver, enabling automatic token exchange for external identity provider tokens. The implementation includes a FederationProvider that wraps base token providers and intelligently determines when token exchange is needed by comparing JWT issuers with the Databricks host. It also adds a CachedTokenProvider to optimize token refresh operations.
Key changes:
- Added federation provider with automatic token exchange detection and fallback
- Implemented comprehensive test coverage for federation scenarios including real-world identity providers
- Added caching layer for token providers to reduce unnecessary token refreshes
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| connector.go | Added public API functions for configuring federated token providers with optional client ID support |
| auth/tokenprovider/exchange.go | Implements core federation logic including JWT validation, host comparison, and token exchange protocol |
| auth/tokenprovider/federation_test.go | Comprehensive test suite covering host comparison, token exchange, caching, and real-world identity providers |
| auth/tokenprovider/cached.go | Generic token caching provider with thread-safe refresh logic |
| auth/oauth/oauth.go | Reorganized Azure domain lists, moving staging/dev domains from tenant map to domain list |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| func (p *FederationProvider) Name() string { | ||
| baseName := p.baseProvider.Name() | ||
| if p.clientID != "" { | ||
| return fmt.Sprintf("federation[%s,sp:%s]", baseName, p.clientID[:8]) // Truncate client ID for readability |
There was a problem hiding this comment.
Potential panic if clientID length is less than 8 characters. Add a length check before slicing or use a safe truncation approach.
| return fmt.Sprintf("federation[%s,sp:%s]", baseName, p.clientID[:8]) // Truncate client ID for readability | |
| clientIDDisplay := p.clientID | |
| if len(p.clientID) >= 8 { | |
| clientIDDisplay = p.clientID[:8] | |
| } | |
| return fmt.Sprintf("federation[%s,sp:%s]", baseName, clientIDDisplay) // Truncate client ID for readability |
| fedProvider := NewFederationProviderWithClientID(baseProvider, "test.databricks.com", "client-12345678-more") | ||
| // Should truncate client ID to first 8 chars | ||
| assert.Equal(t, "federation[static,sp:client-1]", fedProvider.Name()) |
There was a problem hiding this comment.
Test assumes client ID will always be at least 8 characters. Add test case for short client IDs (< 8 characters) to verify behavior matches the truncation logic in Name().
Adds token federation for databricks sql go driver