Merge pull request #1 from better/derek/add-host-variable

Add host option to pass down to driver

Author: Dirrk

pkg/provider/provider.go

@@ -131,6 +131,12 @@ func Provider() *schema.Provider {
 				Required:    true,
 				DefaultFunc: schema.EnvDefaultFunc("SNOWFLAKE_REGION", "us-west-2"),
 			},
+			"host": {
+				Type:        schema.TypeString,
+				Description: "Supports passing in a custom host value to the snowflake go driver for use with privatelink",
+				Optional:    true,
+				DefaultFunc: schema.EnvDefaultFunc("SNOWFLAKE_HOST", nil),
+			},
 		},
 		ResourcesMap:   getResources(),
 		DataSourcesMap: getDataSources(),
@@ -259,6 +265,7 @@ func ConfigureProvider(s *schema.ResourceData) (interface{}, error) {
 	oauthClientSecret := s.Get("oauth_client_secret").(string)
 	oauthEndpoint := s.Get("oauth_endpoint").(string)
 	oauthRedirectURL := s.Get("oauth_redirect_url").(string)
+	host := s.Get("host").(string)
 
 	if oauthRefreshToken != "" {
 		accessToken, err := GetOauthAccessToken(oauthEndpoint, oauthClientID, oauthClientSecret, GetOauthData(oauthRefreshToken, oauthRedirectURL))
@@ -279,6 +286,7 @@ func ConfigureProvider(s *schema.ResourceData) (interface{}, error) {
 		oauthAccessToken,
 		region,
 		role,
+		host,
 	)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not build dsn for snowflake connection")
@@ -302,7 +310,8 @@ func DSN(
 	privateKeyPassphrase,
 	oauthAccessToken,
 	region,
-	role string) (string, error) {
+	role,
+	host string) (string, error) {
 
 	// us-west-2 is their default region, but if you actually specify that it won't trigger their default code
 	//  https://github.com/snowflakedb/gosnowflake/blob/52137ce8c32eaf93b0bd22fc5c7297beff339812/dsn.go#L61
@@ -317,6 +326,12 @@ func DSN(
 		Role:    role,
 	}
 
+	// If host is set trust it and do not use the region value
+	if host != "" {
+		config.Region = ""
+		config.Host = host
+	}
+
 	if privateKeyPath != "" {
 		privateKeyBytes, err := ReadPrivateKeyFile(privateKeyPath)
 		if err != nil {

pkg/provider/provider_test.go

@@ -33,22 +33,27 @@ func TestDSN(t *testing.T) {
 		password string
 		browserAuth bool
 		region,
-		role string
+		role,
+		host string
 	}
 	tests := []struct {
 		name    string
 		args    args
 		want    string
 		wantErr bool
 	}{
-		{"simple", args{"acct", "user", "pass", false, "region", "role"},
+		{"simple", args{"acct", "user", "pass", false, "region", "role", ""},
 			"user:pass@acct.region.snowflakecomputing.com:443?ocspFailOpen=true&region=region&role=role&validateDefaultParameters=true", false},
-		{"us-west-2 special case", args{"acct2", "user2", "pass2", false, "us-west-2", "role2"},
+		{"us-west-2 special case", args{"acct2", "user2", "pass2", false, "us-west-2", "role2", ""},
 			"user2:pass2@acct2.snowflakecomputing.com:443?ocspFailOpen=true&role=role2&validateDefaultParameters=true", false},
+		{"customhostwregion", args{"acct3", "user3", "pass3", false, "", "role3", "zha123.us-east-1.privatelink.snowflakecomputing.com"},
+			"user3:pass3@zha123.us-east-1.privatelink.snowflakecomputing.com:443?account=acct3&ocspFailOpen=true&role=role3&validateDefaultParameters=true", false},
+		{"customhostignoreregion", args{"acct4", "user4", "pass4", false, "fakeregion", "role4", "zha1234.us-east-1.privatelink.snowflakecomputing.com"},
+			"user4:pass4@zha1234.us-east-1.privatelink.snowflakecomputing.com:443?account=acct4&ocspFailOpen=true&role=role4&validateDefaultParameters=true", false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := provider.DSN(tt.args.account, tt.args.user, tt.args.password, tt.args.browserAuth, "", "", "", "", tt.args.region, tt.args.role)
+			got, err := provider.DSN(tt.args.account, tt.args.user, tt.args.password, tt.args.browserAuth, "", "", "", "", tt.args.region, tt.args.role, tt.args.host)
 			if (err != nil) != tt.wantErr {
 				t.Errorf("DSN() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -84,7 +89,7 @@ func TestOAuthDSN(t *testing.T) {
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := provider.DSN(tt.args.account, tt.args.user, "", false, "", "", "", tt.args.oauthAccessToken, tt.args.region, tt.args.role)
+			got, err := provider.DSN(tt.args.account, tt.args.user, "", false, "", "", "", tt.args.oauthAccessToken, tt.args.region, tt.args.role, "")
 
 			if (err != nil) != tt.wantErr {
 				t.Errorf("DSN() error = %v, dsn = %v, wantErr %v", err, got, tt.wantErr)