Merge branch 'main' into token-provider-federation

Author: madhav-db

DESIGN_MULTI_STATEMENT_TRANSACTIONS.md

@@ -57,7 +57,7 @@ func (c *authClient) Authenticate(r *http.Request) error {
 
 	c.tokenSource = GetTokenSource(config)
 	token, err := c.tokenSource.Token()
-	log.Info().Msgf("token fetched successfully")
+	log.Debug().Msgf("databricks OAuth token fetched successfully")
 	if err != nil {
 		log.Err(err).Msg("failed to get token")
 

auth/oauth/m2m/m2m.go

@@ -9,7 +9,16 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
-// TokenProviderAuthenticator implements auth.Authenticator using a TokenProvider
+// TokenProviderAuthenticator implements auth.Authenticator using a TokenProvider.
+//
+// Authentication Flow:
+// 1. On each Authenticate() call, retrieves a token from the configured TokenProvider
+// 2. The provider may implement its own caching and refresh logic
+// 3. Validates the returned token is non-empty
+// 4. Sets the Authorization header with the token type and value
+//
+// The authenticator delegates all token management (caching, refresh, expiry)
+// to the underlying TokenProvider implementation.
 type TokenProviderAuthenticator struct {
 	provider TokenProvider
 }

auth/tokenprovider/authenticator.go

@@ -104,32 +104,4 @@ func TestTokenProviderAuthenticator(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, "Bearer external-token-456", req.Header.Get("Authorization"))
 	})
-
-	t.Run("cached_provider_integration", func(t *testing.T) {
-		callCount := 0
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				callCount++
-				return &Token{
-					AccessToken: "cached-token",
-					TokenType:   "Bearer",
-				}, nil
-			},
-			name: "test",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-		authenticator := NewAuthenticator(cachedProvider)
-
-		// Multiple authentication attempts
-		for i := 0; i < 3; i++ {
-			req, _ := http.NewRequest("GET", "http://example.com", nil)
-			err := authenticator.Authenticate(req)
-			require.NoError(t, err)
-			assert.Equal(t, "Bearer cached-token", req.Header.Get("Authorization"))
-		}
-
-		// Should only call base provider once due to caching
-		assert.Equal(t, 1, callCount)
-	})
 }

auth/tokenprovider/authenticator_test.go

@@ -6,41 +6,43 @@ import (
 	"time"
 )
 
-// ExternalTokenProvider provides tokens from an external source (passthrough)
+// ExternalTokenProvider provides tokens from an external source (passthrough).
+// This provider calls a user-supplied function to retrieve tokens on-demand.
 type ExternalTokenProvider struct {
-	tokenFunc func() (string, error)
-	tokenType string
+	tokenSource func() (string, error)
+	tokenType   string
 }
 
 // NewExternalTokenProvider creates a provider that gets tokens from an external function
-func NewExternalTokenProvider(tokenFunc func() (string, error)) *ExternalTokenProvider {
+func NewExternalTokenProvider(tokenSource func() (string, error)) *ExternalTokenProvider {
 	return &ExternalTokenProvider{
-		tokenFunc: tokenFunc,
-		tokenType: "Bearer",
+		tokenSource: tokenSource,
+		tokenType:   "Bearer",
 	}
 }
 
 // NewExternalTokenProviderWithType creates a provider with a custom token type
-func NewExternalTokenProviderWithType(tokenFunc func() (string, error), tokenType string) *ExternalTokenProvider {
+func NewExternalTokenProviderWithType(tokenSource func() (string, error), tokenType string) *ExternalTokenProvider {
 	return &ExternalTokenProvider{
-		tokenFunc: tokenFunc,
-		tokenType: tokenType,
+		tokenSource: tokenSource,
+		tokenType:   tokenType,
 	}
 }
 
 // GetToken retrieves the token from the external source
 func (p *ExternalTokenProvider) GetToken(ctx context.Context) (*Token, error) {
-	if p.tokenFunc == nil {
-		return nil, fmt.Errorf("external token provider: token function is nil")
+	// Check for cancellation first
+	if err := ctx.Err(); err != nil {
+		return nil, fmt.Errorf("external token provider: context cancelled: %w", err)
 	}
 
-	accessToken, err := p.tokenFunc()
-	if err != nil {
-		return nil, fmt.Errorf("external token provider: failed to get token: %w", err)
+	if p.tokenSource == nil {
+		return nil, fmt.Errorf("external token provider: token source is nil")
 	}
 
-	if accessToken == "" {
-		return nil, fmt.Errorf("external token provider: empty token returned")
+	accessToken, err := p.tokenSource()
+	if err != nil {
+		return nil, fmt.Errorf("external token provider: failed to get token: %w", err)
 	}
 
 	return &Token{

auth/tokenprovider/external.go

@@ -29,8 +29,9 @@ func (t *Token) IsExpired() bool {
 	if t.ExpiresAt.IsZero() {
 		return false // No expiry means token doesn't expire
 	}
-	// Consider token expired 5 minutes before actual expiry for safety
-	return time.Now().Add(5 * time.Minute).After(t.ExpiresAt)
+	// Consider token expired 30 seconds before actual expiry for safety
+	// This matches the standard buffer used by other Databricks SDKs
+	return time.Now().Add(30 * time.Second).After(t.ExpiresAt)
 }
 
 // SetAuthHeader sets the Authorization header on an HTTP request

auth/tokenprovider/provider.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"errors"
 	"net/http"
-	"sync"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -44,12 +42,12 @@ func TestToken_IsExpired(t *testing.T) {
 			expected: false,
 		},
 		{
-			name: "token_expires_within_5_minutes",
+			name: "token_expires_within_30_seconds",
 			token: &Token{
 				AccessToken: "test-token",
-				ExpiresAt:   time.Now().Add(3 * time.Minute),
+				ExpiresAt:   time.Now().Add(15 * time.Second),
 			},
-			expected: true, // Should be considered expired due to 5-minute buffer
+			expected: true, // Should be considered expired due to 30-second buffer
 		},
 	}
 
@@ -173,17 +171,18 @@ func TestExternalTokenProvider(t *testing.T) {
 		assert.Contains(t, err.Error(), "failed to get token")
 	})
 
-	t.Run("empty_token_error", func(t *testing.T) {
+	t.Run("empty_token_allowed", func(t *testing.T) {
 		tokenFunc := func() (string, error) {
 			return "", nil
 		}
 
 		provider := NewExternalTokenProvider(tokenFunc)
 		token, err := provider.GetToken(context.Background())
 
-		assert.Error(t, err)
-		assert.Nil(t, token)
-		assert.Contains(t, err.Error(), "empty token returned")
+		assert.NoError(t, err)
+		assert.NotNil(t, token)
+		assert.Empty(t, token.AccessToken)
+		// Empty tokens are validated at the authenticator level, not provider level
 	})
 
 	t.Run("nil_function_error", func(t *testing.T) {
@@ -192,7 +191,7 @@ func TestExternalTokenProvider(t *testing.T) {
 
 		assert.Error(t, err)
 		assert.Nil(t, token)
-		assert.Contains(t, err.Error(), "token function is nil")
+		assert.Contains(t, err.Error(), "token source is nil")
 	})
 
 	t.Run("custom_token_type", func(t *testing.T) {
@@ -228,183 +227,6 @@ func TestExternalTokenProvider(t *testing.T) {
 	})
 }
 
-func TestCachedTokenProvider(t *testing.T) {
-	t.Run("caches_valid_token", func(t *testing.T) {
-		callCount := 0
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				callCount++
-				return &Token{
-					AccessToken: "cached-token",
-					TokenType:   "Bearer",
-					ExpiresAt:   time.Now().Add(1 * time.Hour),
-				}, nil
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-
-		// First call - should fetch from base provider
-		token1, err1 := cachedProvider.GetToken(context.Background())
-		require.NoError(t, err1)
-		assert.Equal(t, "cached-token", token1.AccessToken)
-		assert.Equal(t, 1, callCount)
-
-		// Second call - should use cache
-		token2, err2 := cachedProvider.GetToken(context.Background())
-		require.NoError(t, err2)
-		assert.Equal(t, "cached-token", token2.AccessToken)
-		assert.Equal(t, 1, callCount) // Should still be 1
-	})
-
-	t.Run("refreshes_expired_token", func(t *testing.T) {
-		callCount := 0
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				callCount++
-				// Return token that expires soon
-				return &Token{
-					AccessToken: "token-" + string(rune(callCount)),
-					TokenType:   "Bearer",
-					ExpiresAt:   time.Now().Add(2 * time.Minute), // Within refresh threshold
-				}, nil
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-		cachedProvider.RefreshThreshold = 5 * time.Minute
-
-		// First call
-		token1, err1 := cachedProvider.GetToken(context.Background())
-		require.NoError(t, err1)
-		assert.Equal(t, "token-\x01", token1.AccessToken)
-		assert.Equal(t, 1, callCount)
-
-		// Second call - should refresh because token expires within threshold
-		token2, err2 := cachedProvider.GetToken(context.Background())
-		require.NoError(t, err2)
-		assert.Equal(t, "token-\x02", token2.AccessToken)
-		assert.Equal(t, 2, callCount)
-	})
-
-	t.Run("handles_provider_error", func(t *testing.T) {
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				return nil, errors.New("provider error")
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-		token, err := cachedProvider.GetToken(context.Background())
-
-		assert.Error(t, err)
-		assert.Nil(t, token)
-		assert.Contains(t, err.Error(), "provider error")
-	})
-
-	t.Run("no_expiry_token_not_refreshed", func(t *testing.T) {
-		callCount := 0
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				callCount++
-				return &Token{
-					AccessToken: "permanent-token",
-					TokenType:   "Bearer",
-					ExpiresAt:   time.Time{}, // No expiry
-				}, nil
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-
-		// Multiple calls should all use cache
-		for i := 0; i < 5; i++ {
-			token, err := cachedProvider.GetToken(context.Background())
-			require.NoError(t, err)
-			assert.Equal(t, "permanent-token", token.AccessToken)
-		}
-
-		assert.Equal(t, 1, callCount) // Should only be called once
-	})
-
-	t.Run("clear_cache", func(t *testing.T) {
-		callCount := 0
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				callCount++
-				return &Token{
-					AccessToken: "token-" + string(rune(callCount)),
-					TokenType:   "Bearer",
-					ExpiresAt:   time.Now().Add(1 * time.Hour),
-				}, nil
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-
-		// First call
-		token1, _ := cachedProvider.GetToken(context.Background())
-		assert.Equal(t, "token-\x01", token1.AccessToken)
-		assert.Equal(t, 1, callCount)
-
-		// Clear cache
-		cachedProvider.ClearCache()
-
-		// Next call should fetch new token
-		token2, _ := cachedProvider.GetToken(context.Background())
-		assert.Equal(t, "token-\x02", token2.AccessToken)
-		assert.Equal(t, 2, callCount)
-	})
-
-	t.Run("concurrent_access", func(t *testing.T) {
-		var callCount atomic.Int32
-		baseProvider := &mockProvider{
-			tokenFunc: func() (*Token, error) {
-				// Simulate slow token fetch
-				time.Sleep(100 * time.Millisecond)
-				callCount.Add(1)
-				return &Token{
-					AccessToken: "concurrent-token",
-					TokenType:   "Bearer",
-					ExpiresAt:   time.Now().Add(1 * time.Hour),
-				}, nil
-			},
-			name: "mock",
-		}
-
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-
-		// Launch multiple goroutines
-		var wg sync.WaitGroup
-		for i := 0; i < 10; i++ {
-			wg.Add(1)
-			go func() {
-				defer wg.Done()
-				token, err := cachedProvider.GetToken(context.Background())
-				assert.NoError(t, err)
-				assert.Equal(t, "concurrent-token", token.AccessToken)
-			}()
-		}
-
-		wg.Wait()
-
-		// Should only fetch token once despite concurrent access
-		assert.Equal(t, int32(1), callCount.Load())
-	})
-
-	t.Run("provider_name", func(t *testing.T) {
-		baseProvider := &mockProvider{name: "test-provider"}
-		cachedProvider := NewCachedTokenProvider(baseProvider)
-
-		assert.Equal(t, "cached[test-provider]", cachedProvider.Name())
-	})
-}
-
 // Mock provider for testing
 type mockProvider struct {
 	tokenFunc func() (*Token, error)