feat: add Workload Identity Federation (WIF) authentication support

Adds support for Snowflake Workload Identity Federation as a new
authentication method for the admin connection. Supported providers:
  - AWS  (STS/SigV4, no token config needed)
  - GCP  (metadata service, no token config needed)
  - AZURE (managed identity, optional entra_resource)
  - OIDC (requires workload_identity_token)

New config fields:
  workload_identity_provider       - required, one of AWS/GCP/AZURE/OIDC
  workload_identity_token          - required for OIDC provider only
  workload_identity_entra_resource - optional, Azure only

WIF is mutually exclusive with password and private_key auth.
Adds 21 unit tests covering config construction, connection setup,
and validation across all provider types.

Author: bbhorrigan

connection_producer.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"net/url"
 	"regexp"
+	"strings"
 	"sync"
 	"time"
 
@@ -25,22 +26,38 @@ import (
 	"github.com/snowflakedb/gosnowflake"
 )
 
+const (
+	wifProviderAWS   = "AWS"
+	wifProviderGCP   = "GCP"
+	wifProviderAzure = "AZURE"
+	wifProviderOIDC  = "OIDC"
+)
+
+var validWIFProviders = []string{wifProviderAWS, wifProviderGCP, wifProviderAzure, wifProviderOIDC}
+
 var (
 	ErrInvalidSnowflakeURL           = fmt.Errorf("invalid connection URL format, expect <account_name>.snowflakecomputing.com/<db_name>")
 	ErrInvalidPrivateKey             = fmt.Errorf("failed to read provided private_key")
+	ErrWIFMutuallyExclusive          = fmt.Errorf("workload_identity_provider cannot be combined with password or private_key authentication")
+	ErrWIFTokenRequired              = fmt.Errorf("workload_identity_token is required when using the OIDC workload identity provider")
+	ErrWIFUsernameRequired           = fmt.Errorf("username is required for workload identity federation")
+	ErrWIFInvalidProvider            = fmt.Errorf("workload_identity_provider must be one of: AWS, GCP, AZURE, OIDC")
 	accountAndDBNameFromConnURLRegex = regexp.MustCompile(`^(.+)\.snowflakecomputing\.com/(.+)$`) // Expected format: <account_name>.snowflakecomputing.com/<db_name>
 )
 
 type snowflakeConnectionProducer struct {
-	ConnectionURL            string      `json:"connection_url"`
-	MaxOpenConnections       int         `json:"max_open_connections"`
-	MaxIdleConnections       int         `json:"max_idle_connections"`
-	MaxConnectionLifetimeRaw interface{} `json:"max_connection_lifetime"`
-	Username                 string      `json:"username"`
-	Password                 string      `json:"password"`
-	PrivateKey               []byte      `json:"private_key"`
-	UsernameTemplate         string      `json:"username_template"`
-	DisableEscaping          bool        `json:"disable_escaping"`
+	ConnectionURL                string      `json:"connection_url"`
+	MaxOpenConnections           int         `json:"max_open_connections"`
+	MaxIdleConnections           int         `json:"max_idle_connections"`
+	MaxConnectionLifetimeRaw     interface{} `json:"max_connection_lifetime"`
+	Username                     string      `json:"username"`
+	Password                     string      `json:"password"`
+	PrivateKey                   []byte      `json:"private_key"`
+	WorkloadIdentityProvider     string      `json:"workload_identity_provider"`
+	WorkloadIdentityToken        string      `json:"workload_identity_token"`
+	WorkloadIdentityEntraResource string     `json:"workload_identity_entra_resource"`
+	UsernameTemplate             string      `json:"username_template"`
+	DisableEscaping              bool        `json:"disable_escaping"`
 
 	Initialized           bool
 	RawConfig             map[string]any
@@ -53,8 +70,9 @@ type snowflakeConnectionProducer struct {
 
 func (c *snowflakeConnectionProducer) secretValues() map[string]string {
 	return map[string]string{
-		c.Password:           "[password]",
-		string(c.PrivateKey): "[private_key]",
+		c.Password:                "[password]",
+		string(c.PrivateKey):      "[private_key]",
+		c.WorkloadIdentityToken:   "[workload_identity_token]",
 	}
 }
 
@@ -86,7 +104,11 @@ func (c *snowflakeConnectionProducer) Init(ctx context.Context, initConfig map[s
 		return nil, fmt.Errorf("connection_url cannot be empty")
 	}
 
-	if len(c.Password) > 0 {
+	if len(c.WorkloadIdentityProvider) > 0 {
+		if err := c.validateWIFConfig(); err != nil {
+			return nil, err
+		}
+	} else if len(c.Password) > 0 {
 		// Return an error here once Snowflake ends support for password auth.
 		c.logger.Warn("[DEPRECATED] Single-factor password authentication is deprecated in Snowflake and will be removed by November 2025. " +
 			"Key pair authentication will be required after this date.")
@@ -137,6 +159,36 @@ func (c *snowflakeConnectionProducer) Init(ctx context.Context, initConfig map[s
 	return initConfig, nil
 }
 
+// validateWIFConfig checks that the WIF configuration is valid and mutually
+// exclusive with other authentication methods.
+func (c *snowflakeConnectionProducer) validateWIFConfig() error {
+	if len(c.Password) > 0 || len(c.PrivateKey) > 0 {
+		return ErrWIFMutuallyExclusive
+	}
+
+	if c.Username == "" {
+		return ErrWIFUsernameRequired
+	}
+
+	provider := strings.ToUpper(c.WorkloadIdentityProvider)
+	valid := false
+	for _, p := range validWIFProviders {
+		if provider == p {
+			valid = true
+			break
+		}
+	}
+	if !valid {
+		return ErrWIFInvalidProvider
+	}
+
+	if provider == wifProviderOIDC && c.WorkloadIdentityToken == "" {
+		return ErrWIFTokenRequired
+	}
+
+	return nil
+}
+
 func (c *snowflakeConnectionProducer) Initialize(ctx context.Context, config map[string]any, verifyConnection bool) error {
 	_, err := c.Init(ctx, config, verifyConnection)
 	return err
@@ -162,12 +214,18 @@ func (c *snowflakeConnectionProducer) Connection(ctx context.Context) (interface
 
 	var db *sql.DB
 	var err error
-	if len(c.PrivateKey) > 0 {
+	switch {
+	case len(c.WorkloadIdentityProvider) > 0:
+		db, err = openSnowflakeWIF(c.ConnectionURL, c.Username, c.WorkloadIdentityProvider, c.WorkloadIdentityToken, c.WorkloadIdentityEntraResource)
+		if err != nil {
+			return nil, fmt.Errorf("error opening Snowflake connection using workload identity federation: %w", err)
+		}
+	case len(c.PrivateKey) > 0:
 		db, err = openSnowflake(c.ConnectionURL, c.Username, c.PrivateKey)
 		if err != nil {
 			return nil, fmt.Errorf("error opening Snowflake connection using key-pair auth: %w", err)
 		}
-	} else {
+	default:
 		db, err = sql.Open(snowflakeSQLTypeName, c.ConnectionURL)
 		if err != nil {
 			return nil, fmt.Errorf("error opening Snowflake connection using user-pass auth: %w", err)
@@ -202,6 +260,50 @@ func (c *snowflakeConnectionProducer) Close() error {
 	return c.close()
 }
 
+// openSnowflakeWIF opens a Snowflake connection using Workload Identity Federation.
+// provider must be one of: AWS, GCP, AZURE, OIDC.
+// token is required only for the OIDC provider.
+// entraResource is optional and only applies to Azure environments.
+func openSnowflakeWIF(connectionURL, username, provider, token, entraResource string) (*sql.DB, error) {
+	cfg, err := getSnowflakeWIFConfig(connectionURL, username, provider, token, entraResource)
+	if err != nil {
+		return nil, fmt.Errorf("error constructing snowflake WIF config: %w", err)
+	}
+	connector := gosnowflake.NewConnector(gosnowflake.SnowflakeDriver{}, *cfg)
+	return sql.OpenDB(connector), nil
+}
+
+// getSnowflakeWIFConfig builds a gosnowflake.Config for Workload Identity Federation auth.
+func getSnowflakeWIFConfig(connectionURL, username, provider, token, entraResource string) (*gosnowflake.Config, error) {
+	u, err := url.Parse(connectionURL)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing Snowflake connection URL %q: %w", connectionURL, err)
+	}
+
+	q := u.Query()
+	q.Set("authenticator", gosnowflake.AuthTypeWorkloadIdentityFederation.String())
+	u.RawQuery = q.Encode()
+
+	// construct dsn: "user:@<account>.snowflakecomputing.com/<db>?..."
+	dsn := fmt.Sprintf("%s:%s@%s", username, "", u.String())
+	cfg, err := gosnowflake.ParseDSN(dsn)
+	if err != nil {
+		return nil, fmt.Errorf("error parsing Snowflake DSN: %w", err)
+	}
+
+	cfg.WorkloadIdentityProvider = strings.ToUpper(provider)
+
+	if token != "" {
+		cfg.Token = token
+	}
+
+	if entraResource != "" {
+		cfg.WorkloadIdentityEntraResource = entraResource
+	}
+
+	return cfg, nil
+}
+
 // Open the DB connection to Snowflake or return an error.
 func openSnowflake(connectionURL, username string, providedPrivateKey []byte) (*sql.DB, error) {
 	cfg, err := getSnowflakeConfig(connectionURL, username, providedPrivateKey)