Consolidate serviceToServiceVisitor and serviceToServiceVisitorWithFallback into one function

Author: hectorcast-db

config/auth_azure_cli.go

@@ -52,7 +52,7 @@ func (c AzureCliCredentials) getVisitor(ctx context.Context, cfg *Config, inner
 		return azureVisitor(cfg, refreshableVisitor(inner, opts...)), nil
 	}
 	management := azureReuseTokenSource(t, ts, opts...)
-	return azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, opts...)), nil
+	return azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, false, opts...)), nil
 }
 
 func (c AzureCliCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {

config/auth_azure_client_secret.go

@@ -57,6 +57,6 @@ func (c AzureClientSecretCredentials) Configure(ctx context.Context, cfg *Config
 	opts := cacheOptions(cfg)
 	inner := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, env.AzureApplicationID), opts...)
 	management := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, aadEndpoint, managementEndpoint), opts...)
-	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, opts...))
+	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, false, opts...))
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }

config/auth_azure_msi.go

@@ -46,7 +46,7 @@ func (c AzureMsiCredentials) Configure(ctx context.Context, cfg *Config) (creden
 	opts := cacheOptions(cfg)
 	inner := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, "", env.AzureApplicationID), opts...)
 	management := azureReuseTokenSource(nil, c.tokenSourceFor(ctx, cfg, "", env.AzureServiceManagementEndpoint()), opts...)
-	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, opts...))
+	visitor := azureVisitor(cfg, serviceToServiceVisitor(inner, management, xDatabricksAzureSpManagementToken, false, opts...))
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }
 

config/auth_gcp_google_credentials.go

@@ -42,7 +42,7 @@ func (c GoogleCredentials) Configure(ctx context.Context, cfg *Config) (credenti
 		return nil, fmt.Errorf("could not obtain OAuth2 token from JSON: %w", err)
 	}
 	logger.Infof(ctx, "Using Google Credentials")
-	visitor := serviceToServiceVisitorWithFallback(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token", cacheOptions(cfg)...)
+	visitor := serviceToServiceVisitor(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token", true, cacheOptions(cfg)...)
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }
 

config/auth_gcp_google_id.go

@@ -44,7 +44,7 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 		return credentials.CredentialsProviderFn(visitor), nil
 	}
 	logger.Infof(ctx, "Using Google Default Application Credentials")
-	visitor := serviceToServiceVisitorWithFallback(inner, platform, "X-Databricks-GCP-SA-Access-Token", opts...)
+	visitor := serviceToServiceVisitor(inner, platform, "X-Databricks-GCP-SA-Access-Token", true, opts...)
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }
 

config/oauth_visitors.go

@@ -20,9 +20,11 @@ func cacheOptions(cfg *Config) []auth.Option {
 }
 
 // serviceToServiceVisitor returns a visitor that sets the Authorization header
-// to the token from the auth token source and the provided secondary header to
-// the token from the secondary token source.
-func serviceToServiceVisitor(primary, secondary oauth2.TokenSource, secondaryHeader string, opts ...auth.Option) func(r *http.Request) error {
+// to the token from the primary token source and the provided secondary header
+// to the token from the secondary token source. If secondaryOptional is true,
+// a failure to get the secondary token logs a warning and skips the header
+// instead of returning an error.
+func serviceToServiceVisitor(primary, secondary oauth2.TokenSource, secondaryHeader string, secondaryOptional bool, opts ...auth.Option) func(r *http.Request) error {
 	refreshableAuth := auth.NewCachedTokenSource(authconv.AuthTokenSource(primary), opts...)
 	refreshableSecondary := auth.NewCachedTokenSource(authconv.AuthTokenSource(secondary), opts...)
 	return func(r *http.Request) error {
@@ -34,36 +36,17 @@ func serviceToServiceVisitor(primary, secondary oauth2.TokenSource, secondaryHea
 
 		cloud, err := refreshableSecondary.Token(context.Background())
 		if err != nil {
+			if secondaryOptional {
+				logger.Warnf(r.Context(), "Failed to get secondary token for %s header: %v. Skipping.", secondaryHeader, err)
+				return nil
+			}
 			return fmt.Errorf("cloud token: %w", err)
 		}
 		r.Header.Set(secondaryHeader, cloud.AccessToken)
 		return nil
 	}
 }
 
-// serviceToServiceVisitorWithFallback is like serviceToServiceVisitor but
-// logs a warning and skips the secondary header when the secondary token
-// source fails, instead of returning an error.
-func serviceToServiceVisitorWithFallback(primary, secondary oauth2.TokenSource, secondaryHeader string, opts ...auth.Option) func(r *http.Request) error {
-	refreshableAuth := auth.NewCachedTokenSource(authconv.AuthTokenSource(primary), opts...)
-	refreshableSecondary := auth.NewCachedTokenSource(authconv.AuthTokenSource(secondary), opts...)
-	return func(r *http.Request) error {
-		inner, err := refreshableAuth.Token(context.Background())
-		if err != nil {
-			return fmt.Errorf("inner token: %w", err)
-		}
-		inner.SetAuthHeader(r)
-
-		cloud, err := refreshableSecondary.Token(context.Background())
-		if err != nil {
-			logger.Warnf(r.Context(), "Failed to get secondary token for %s header: %v. Skipping.", secondaryHeader, err)
-			return nil
-		}
-		r.Header.Set(secondaryHeader, cloud.AccessToken)
-		return nil
-	}
-}
-
 // The same as serviceToServiceVisitor, but without a secondary token source.
 func refreshableVisitor(inner oauth2.TokenSource, opts ...auth.Option) func(r *http.Request) error {
 	return refreshableAuthVisitor(authconv.AuthTokenSource(inner), opts...)

config/oauth_visitors_test.go

@@ -47,7 +47,7 @@ func (s *staticTokenSource) Token() (*oauth2.Token, error) {
 func TestServiceToServiceVisitorWithFallback_BothSucceed(t *testing.T) {
 	primary := &staticTokenSource{token: &oauth2.Token{AccessToken: "primary-token"}}
 	secondary := &staticTokenSource{token: &oauth2.Token{AccessToken: "secondary-token"}}
-	visitor := serviceToServiceVisitorWithFallback(primary, secondary, "X-Secondary")
+	visitor := serviceToServiceVisitor(primary, secondary, "X-Secondary", true)
 
 	req, err := http.NewRequest("GET", "https://example.com", nil)
 	require.NoError(t, err)
@@ -60,7 +60,7 @@ func TestServiceToServiceVisitorWithFallback_BothSucceed(t *testing.T) {
 func TestServiceToServiceVisitorWithFallback_SecondaryFails_SkipsHeader(t *testing.T) {
 	primary := &staticTokenSource{token: &oauth2.Token{AccessToken: "primary-token"}}
 	secondary := &staticTokenSource{err: fmt.Errorf("secondary failed")}
-	visitor := serviceToServiceVisitorWithFallback(primary, secondary, "X-Secondary")
+	visitor := serviceToServiceVisitor(primary, secondary, "X-Secondary", true)
 
 	req, err := http.NewRequest("GET", "https://example.com", nil)
 	require.NoError(t, err)
@@ -73,7 +73,7 @@ func TestServiceToServiceVisitorWithFallback_SecondaryFails_SkipsHeader(t *testi
 func TestServiceToServiceVisitorWithFallback_PrimaryFails_ReturnsError(t *testing.T) {
 	primary := &staticTokenSource{err: fmt.Errorf("primary failed")}
 	secondary := &staticTokenSource{token: &oauth2.Token{AccessToken: "secondary-token"}}
-	visitor := serviceToServiceVisitorWithFallback(primary, secondary, "X-Secondary")
+	visitor := serviceToServiceVisitor(primary, secondary, "X-Secondary", true)
 
 	req, err := http.NewRequest("GET", "https://example.com", nil)
 	require.NoError(t, err)