Updated snowflake config for OAuth

rudderlabs · Dec 16, 2024 · 48b0a02 · 48b0a02
1 parent 78e3a6a
commit 48b0a02
Showing 1 changed file with 56 additions and 34 deletions.
diff --git a/sqlconnect/internal/snowflake/config.go b/sqlconnect/internal/snowflake/config.go
@@ -20,6 +20,7 @@ type Config struct {
 	User      string `json:"user"`
 	Schema    string `json:"schema"`
 	Role      string `json:"role"`
+	Region    string `json:"region"`
 
 	Password string `json:"password"`
 
@@ -40,45 +41,66 @@ type Config struct {
 }
 
 func (c Config) ConnectionString() (dsn string, err error) {
-	sc := gosnowflake.Config{
-		Authenticator: gosnowflake.AuthTypeSnowflake,
-		User:          c.User,
-		Password:      c.Password,
-		Account:       c.Account,
-		Database:      c.DBName,
-		Warehouse:     c.Warehouse,
-		Schema:        c.Schema,
-		Role:          c.Role,
-		Application:   c.Application,
-		LoginTimeout:  c.LoginTimeout,
-		Params:        make(map[string]*string),
-	}
-
-	if c.UseKeyPairAuth {
-		sc.Authenticator = gosnowflake.AuthTypeJwt
-		privateKey, err := c.ParsePrivateKey()
+	if c.UseOAuth {
+		sc := gosnowflake.Config{
+			Authenticator:    gosnowflake.AuthTypeOAuth,
+			Account:          c.Account,
+			Region:           c.Region,
+			Token:            c.OAuthToken,
+			Warehouse:        c.Warehouse,
+			Schema:           c.Schema,
+			Database:         c.DBName,
+			Host:             c.Host,
+			Protocol:         "https",
+			Port:             443,
+			KeepSessionAlive: true,
+		}
+		dsn, err = gosnowflake.DSN(&sc)
 		if err != nil {
-			return "", fmt.Errorf("parsing private key: %w", err)
+			err = fmt.Errorf("creating dsn: %v", err)
+		}
+	} else {
+		sc := gosnowflake.Config{
+			Authenticator: gosnowflake.AuthTypeSnowflake,
+			User:          c.User,
+			Password:      c.Password,
+			Account:       c.Account,
+			Database:      c.DBName,
+			Warehouse:     c.Warehouse,
+			Schema:        c.Schema,
+			Role:          c.Role,
+			Application:   c.Application,
+			LoginTimeout:  c.LoginTimeout,
+			Params:        make(map[string]*string),
 		}
-		sc.PrivateKey = privateKey
-	} else if c.UseOAuth {
-		sc.Authenticator = gosnowflake.AuthTypeOAuth
-		sc.Host = c.Host
-		sc.Token = c.OAuthToken
-	}
 
-	if c.KeepSessionAlive {
-		valueTrue := "true"
-		sc.Params["client_session_keep_alive"] = &valueTrue
-	}
+		if c.UseKeyPairAuth {
+			sc.Authenticator = gosnowflake.AuthTypeJwt
+			privateKey, err := c.ParsePrivateKey()
+			if err != nil {
+				return "", fmt.Errorf("parsing private key: %w", err)
+			}
+			sc.PrivateKey = privateKey
+		} else if c.UseOAuth {
+			sc.Authenticator = gosnowflake.AuthTypeOAuth
+			sc.Host = c.Host
+			sc.Token = c.OAuthToken
+			sc.User = c.User
+		}
 
-	if c.QueryTag != "" {
-		sc.Params["query_tag"] = &c.QueryTag
-	}
+		if c.KeepSessionAlive {
+			valueTrue := "true"
+			sc.Params["client_session_keep_alive"] = &valueTrue
+		}
 
-	dsn, err = gosnowflake.DSN(&sc)
-	if err != nil {
-		err = fmt.Errorf("creating dsn: %v", err)
+		if c.QueryTag != "" {
+			sc.Params["query_tag"] = &c.QueryTag
+		}
+
+		dsn, err = gosnowflake.DSN(&sc)
+		if err != nil {
+			err = fmt.Errorf("creating dsn: %v", err)
+		}
 	}
 	return
 }