From da9cada6b5ea319ff6b5ac1bb59cb619f11cc008 Mon Sep 17 00:00:00 2001 From: Antoine Niek Date: Tue, 22 Feb 2022 15:44:44 -0500 Subject: [PATCH 1/2] Support DIRECTORY option on stage management --- docs/resources/stage.md | 1 + pkg/resources/stage.go | 10 ++++++++++ pkg/snowflake/stage.go | 11 +++++++++++ pkg/snowflake/stage_test.go | 7 +++++-- 4 files changed, 27 insertions(+), 2 deletions(-) diff --git a/docs/resources/stage.md b/docs/resources/stage.md index 0ed3b82460..bc8c1b7947 100644 --- a/docs/resources/stage.md +++ b/docs/resources/stage.md @@ -45,6 +45,7 @@ resource "snowflake_stage_grant" "grant_example_stage" { - **comment** (String) Specifies a comment for the stage. - **copy_options** (String) Specifies the copy options for the stage. - **credentials** (String, Sensitive) Specifies the credentials for the stage. +- **directory** (String) Specifies the directory settings for the stage. - **encryption** (String) Specifies the encryption settings for the stage. - **file_format** (String) Specifies the file format for the stage. - **id** (String) The ID of this resource. diff --git a/pkg/resources/stage.go b/pkg/resources/stage.go index bc7667b91e..121006c539 100644 --- a/pkg/resources/stage.go +++ b/pkg/resources/stage.go @@ -72,6 +72,12 @@ var stageSchema = map[string]*schema.Schema{ Optional: true, Description: "Specifies a comment for the stage.", }, + "directory": { + Type: schema.TypeString, + ForceNew: true, + Optional: true, + Description: "Specifies the directory settings for the stage.", + }, "aws_external_id": { Type: schema.TypeString, Optional: true, @@ -176,6 +182,10 @@ func CreateStage(d *schema.ResourceData, meta interface{}) error { builder.WithCopyOptions(v.(string)) } + if v, ok := d.GetOk("directory"); ok { + builder.WithDirectory(v.(string)) + } + if v, ok := d.GetOk("encryption"); ok { builder.WithEncryption(v.(string)) } diff --git a/pkg/snowflake/stage.go b/pkg/snowflake/stage.go index 1b12d75ec3..5f5eee1591 100644 --- a/pkg/snowflake/stage.go +++ b/pkg/snowflake/stage.go @@ -21,6 +21,7 @@ type StageBuilder struct { schema string url string credentials string + directory string storageIntegration string encryption string fileFormat string @@ -74,6 +75,12 @@ func (sb *StageBuilder) WithCopyOptions(c string) *StageBuilder { return sb } +// WithDirectory adds directory option to the StageBuilder +func (sb *StageBuilder) WithDirectory(d string) *StageBuilder { + sb.directory = d + return sb +} + // WithComment adds a comment to the StageBuilder func (sb *StageBuilder) WithComment(c string) *StageBuilder { sb.comment = c @@ -150,6 +157,10 @@ func (sb *StageBuilder) Create() string { q.WriteString(fmt.Sprintf(` COPY_OPTIONS = (%v)`, sb.copyOptions)) } + if sb.directory != "" { + q.WriteString(fmt.Sprintf(` DIRECTORY = (%v)`, sb.directory)) + } + if sb.comment != "" { q.WriteString(fmt.Sprintf(` COMMENT = '%v'`, EscapeString(sb.comment))) } diff --git a/pkg/snowflake/stage_test.go b/pkg/snowflake/stage_test.go index dbb3a78797..6a460aec82 100644 --- a/pkg/snowflake/stage_test.go +++ b/pkg/snowflake/stage_test.go @@ -28,11 +28,14 @@ func TestStageCreate(t *testing.T) { s.WithCopyOptions("on_error='skip_file'") r.Equal(s.Create(), `CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file')`) + s.WithDirectory("ENABLE=TRUE") + r.Equal(s.Create(), `CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file') DIRECTORY = (ENABLE=TRUE)`) + s.WithComment("Yee'haw") - r.Equal(`CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file') COMMENT = 'Yee\'haw'`, s.Create()) + r.Equal(`CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file') DIRECTORY = (ENABLE=TRUE) COMMENT = 'Yee\'haw'`, s.Create()) s.WithStorageIntegration("MY_INTEGRATION") - r.Equal(`CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') STORAGE_INTEGRATION = MY_INTEGRATION ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file') COMMENT = 'Yee\'haw'`, s.Create()) + r.Equal(`CREATE STAGE "test_db"."test_schema"."test_stage" URL = 's3://load/encrypted_files/' CREDENTIALS = (aws_role='arn:aws:iam::001234567890:role/mysnowflakerole') STORAGE_INTEGRATION = MY_INTEGRATION ENCRYPTION = (type='AWS_SSE_KMS' kms_key_id = 'aws/key') FILE_FORMAT = (format_name=my_csv_format) COPY_OPTIONS = (on_error='skip_file') DIRECTORY = (ENABLE=TRUE) COMMENT = 'Yee\'haw'`, s.Create()) } func TestStageRename(t *testing.T) { From 054353bd0c2aba4b332a7eed359b1456d4e0e8f0 Mon Sep 17 00:00:00 2001 From: Antoine Niek Date: Tue, 22 Feb 2022 16:13:17 -0500 Subject: [PATCH 2/2] Add read support for directory stage option --- pkg/resources/stage.go | 5 +++++ pkg/resources/stage_test.go | 3 ++- pkg/snowflake/stage.go | 7 +++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/pkg/resources/stage.go b/pkg/resources/stage.go index 121006c539..3589e64a8a 100644 --- a/pkg/resources/stage.go +++ b/pkg/resources/stage.go @@ -283,6 +283,11 @@ func ReadStage(d *schema.ResourceData, meta interface{}) error { return err } + err = d.Set("directory", stageDesc.Directory) + if err != nil { + return err + } + err = d.Set("storage_integration", s.StorageIntegration) if err != nil { return err diff --git a/pkg/resources/stage_test.go b/pkg/resources/stage_test.go index 5d351aac40..1d23a5e68d 100644 --- a/pkg/resources/stage_test.go +++ b/pkg/resources/stage_test.go @@ -73,7 +73,8 @@ func expectReadStage(mock sqlmock.Sqlmock) { "parent_property", "property", "property_type", "property_value", "property_default"}, ).AddRow("STAGE_LOCATION", "URL", "string", `["s3://load/test/"]`, ""). AddRow("STAGE_CREDENTIALS", "AWS_EXTERNAL_ID", "string", "test", ""). - AddRow("STAGE_FILE_FORMAT", "FORMAT_NAME", "string", "CSV", "") + AddRow("STAGE_FILE_FORMAT", "FORMAT_NAME", "string", "CSV", ""). + AddRow("DIRECTORY", "ENABLED", "Boolean", true, false) mock.ExpectQuery(`^DESCRIBE STAGE "test_db"."test_schema"."test_stage"$`).WillReturnRows(rows) } diff --git a/pkg/snowflake/stage.go b/pkg/snowflake/stage.go index 5f5eee1591..4c445a1839 100644 --- a/pkg/snowflake/stage.go +++ b/pkg/snowflake/stage.go @@ -253,6 +253,7 @@ type descStageResult struct { SnowflakeIamUser string FileFormat string CopyOptions string + Directory string } type descStageRow struct { @@ -266,6 +267,7 @@ func DescStage(db *sql.DB, query string) (*descStageResult, error) { r := &descStageResult{} var ff []string var co []string + var dir []string rows, err := Query(db, query) if err != nil { return r, err @@ -297,11 +299,16 @@ func DescStage(db *sql.DB, query string) (*descStageResult, error) { if row.PropertyValue != row.PropertyDefault { co = append(co, fmt.Sprintf("%s = %s", row.Property, row.PropertyValue)) } + case "DIRECTORY": + if row.PropertyValue != row.PropertyDefault { + dir = append(dir, fmt.Sprintf("%s = %s", row.Property, row.PropertyValue)) + } } } r.FileFormat = strings.Join(ff, " ") r.CopyOptions = strings.Join(co, " ") + r.Directory = strings.Join(dir, " ") return r, nil }