feat: S3GOV support to storage_integration (#1133)

* Updating api_provider list * api acceptance test * adding test for aws_gov_api_gateway * changing the region from us-gov-west-2 to us-gov-west-1 in expectRead func * adding S3GOV support to storage integration * adding tests for storage integration with S3GOV as storage provider * reformating Co-authored-by: Jason Lin <[email protected]>
Snowflake-Labs · Jul 14, 2022 · 92a5e35 · 92a5e35
1 parent 7a27b40
commit 92a5e35
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 4 deletions.
diff --git a/pkg/resources/storage_integration.go b/pkg/resources/storage_integration.go
@@ -52,7 +52,7 @@ var storageIntegrationSchema = map[string]*schema.Schema{
 	"storage_provider": {
 		Type:         schema.TypeString,
 		Required:     true,
-		ValidateFunc: validation.StringInSlice([]string{"S3", "GCS", "AZURE"}, false),
+		ValidateFunc: validation.StringInSlice([]string{"S3", "GCS", "AZURE", "S3GOV"}, false),
 	},
 	"storage_aws_external_id": {
 		Type:        schema.TypeString,
@@ -374,7 +374,7 @@ func setStorageProviderSettings(data *schema.ResourceData, stmt snowflake.Settin
 	stmt.SetString("STORAGE_PROVIDER", storageProvider)
 
 	switch storageProvider {
-	case "S3":
+	case "S3", "S3GOV":
 		v, ok := data.GetOk("storage_aws_role_arn")
 		if !ok {
 			return fmt.Errorf("If you use the S3 storage provider you must specify a storage_aws_role_arn")

diff --git a/pkg/resources/storage_integration_test.go b/pkg/resources/storage_integration_test.go
@@ -29,8 +29,21 @@ func TestStorageIntegrationCreate(t *testing.T) {
 		"storage_aws_role_arn":      "we-should-probably-validate-this-string",
 		"storage_aws_object_acl":    "bucket-owner-full-control",
 	}
+
+	in2 := map[string]interface{}{
+		"name":                      "test_storage_integration_with_s3gov",
+		"comment":                   "great comment",
+		"storage_allowed_locations": []interface{}{"s3://great-bucket/great-path/"},
+		"storage_provider":          "S3GOV",
+		"storage_aws_role_arn":      "we-should-probably-validate-this-string",
+		"storage_aws_object_acl":    "bucket-owner-full-control",
+	}
+
 	d := schema.TestResourceDataRaw(t, resources.StorageIntegration().Schema, in)
+	d2 := schema.TestResourceDataRaw(t, resources.StorageIntegration().Schema, in2)
+
 	r.NotNil(d)
+	r.NotNil(d2)
 
 	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
 		mock.ExpectExec(
@@ -41,6 +54,16 @@ func TestStorageIntegrationCreate(t *testing.T) {
 		err := resources.CreateStorageIntegration(d, db)
 		r.NoError(err)
 	})
+
+	WithMockDb(t, func(db *sql.DB, mock sqlmock.Sqlmock) {
+		mock.ExpectExec(
+			`^CREATE STORAGE INTEGRATION "test_storage_integration_with_s3gov" COMMENT='great comment' STORAGE_AWS_OBJECT_ACL='bucket-owner-full-control' STORAGE_AWS_ROLE_ARN='we-should-probably-validate-this-string' STORAGE_PROVIDER='S3GOV' TYPE='EXTERNAL_STAGE' STORAGE_ALLOWED_LOCATIONS=\('s3://great-bucket/great-path/'\) ENABLED=true$`,
+		).WillReturnResult(sqlmock.NewResult(1, 1))
+		expectReadStorageIntegrationWithS3GOV(mock)
+
+		err := resources.CreateStorageIntegration(d2, db)
+		r.NoError(err)
+	})
 }
 
 func TestStorageIntegrationRead(t *testing.T) {
@@ -118,8 +141,8 @@ func expectReadStorageIntegration(mock sqlmock.Sqlmock) {
 	mock.ExpectQuery(`^SHOW STORAGE INTEGRATIONS LIKE 'test_storage_integration'$`).WillReturnRows(showRows)
 
 	descRows := sqlmock.NewRows([]string{
-		"property", "property_type", "property_value", "property_default",
-	}).AddRow("ENABLED", "Boolean", true, false).
+		"property", "property_type", "property_value", "property_default"}).
+		AddRow("ENABLED", "Boolean", true, false).
 		AddRow("STORAGE_PROVIDER", "String", "S3", nil).
 		AddRow("STORAGE_ALLOWED_LOCATIONS", "List", "s3://bucket-a/path-a/,s3://bucket-b/", nil).
 		AddRow("STORAGE_BLOCKED_LOCATIONS", "List", "s3://bucket-c/path-c/,s3://bucket-d/", nil).
@@ -131,6 +154,26 @@ func expectReadStorageIntegration(mock sqlmock.Sqlmock) {
 	mock.ExpectQuery(`DESCRIBE STORAGE INTEGRATION "test_storage_integration"$`).WillReturnRows(descRows)
 }
 
+func expectReadStorageIntegrationWithS3GOV(mock sqlmock.Sqlmock) {
+	showRows := sqlmock.NewRows([]string{
+		"name", "type", "category", "enabled", "created_on"},
+	).AddRow("test_storage_integration_with_s3gov", "EXTERNAL_STAGE", "STORAGE", true, "now")
+	mock.ExpectQuery(`^SHOW STORAGE INTEGRATIONS LIKE 'test_storage_integration_with_s3gov'$`).WillReturnRows(showRows)
+
+	descRows := sqlmock.NewRows([]string{
+		"property", "property_type", "property_value", "property_default"}).
+		AddRow("ENABLED", "Boolean", true, false).
+		AddRow("STORAGE_PROVIDER", "String", "S3GOV", nil).
+		AddRow("STORAGE_ALLOWED_LOCATIONS", "List", "s3://bucket-a/path-a/,s3://bucket-b/", nil).
+		AddRow("STORAGE_BLOCKED_LOCATIONS", "List", "s3://bucket-c/path-c/,s3://bucket-d/", nil).
+		AddRow("STORAGE_AWS_IAM_USER_ARN", "String", "arn:aws:iam::000000000000:/user/test", nil).
+		AddRow("STORAGE_AWS_ROLE_ARN", "String", "arn:aws:iam::000000000001:/role/test", nil).
+		AddRow("STORAGE_AWS_OBJECT_ACL", "String", "bucket-owner-full-control", nil).
+		AddRow("STORAGE_AWS_EXTERNAL_ID", "String", "AGreatExternalID", nil)
+
+	mock.ExpectQuery(`DESCRIBE STORAGE INTEGRATION "test_storage_integration_with_s3gov"$`).WillReturnRows(descRows)
+}
+
 func expectReadStorageIntegrationEmpty(mock sqlmock.Sqlmock) {
 	noRows := sqlmock.NewRows([]string{
 		"name", "type", "category", "enabled", "created_on"},