Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-jcieslak
sfc-gh-jcieslak committed Nov 21, 2024
1 parent c09da00 commit c35bf68
Show file tree
Hide file tree
Showing 9 changed files with 130 additions and 98 deletions.
8 changes: 4 additions & 4 deletions pkg/acceptance/helpers/storage_integration_client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,10 @@ func (c *StorageIntegrationClient) CreateS3(t *testing.T, awsBucketUrl, awsRoleA

id := c.ids.RandomAccountObjectIdentifier()
req := sdk.NewCreateStorageIntegrationRequest(id, true, s3AllowedLocations).
WithIfNotExists(sdk.Bool(true)).
WithS3StorageProviderParams(sdk.NewS3StorageParamsRequest(awsRoleArn)).
WithIfNotExists(true).
WithS3StorageProviderParams(*sdk.NewS3StorageParamsRequest(sdk.RegularS3Protocol, awsRoleArn)).
WithStorageBlockedLocations(s3BlockedLocations).
WithComment(sdk.String("some comment"))
WithComment("some comment")

err := c.client().Create(ctx, req)
require.NoError(t, err)
Expand All @@ -73,7 +73,7 @@ func (c *StorageIntegrationClient) DropFunc(t *testing.T, id sdk.AccountObjectId
ctx := context.Background()

return func() {
err := c.client().Drop(ctx, sdk.NewDropStorageIntegrationRequest(id).WithIfExists(sdk.Bool(true)))
err := c.client().Drop(ctx, sdk.NewDropStorageIntegrationRequest(id).WithIfExists(true))
require.NoError(t, err)
}
}
51 changes: 28 additions & 23 deletions pkg/resources/storage_integration.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"fmt"
"log"
"slices"
"strings"

"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
Expand Down Expand Up @@ -52,12 +53,11 @@ var storageIntegrationSchema = map[string]*schema.Schema{
Optional: true,
Description: "Explicitly prohibits external stages that use the integration from referencing one or more storage locations.",
},
// TODO (SNOW-1015282): Remove S3gov option before going into V1
"storage_provider": {
Type: schema.TypeString,
Required: true,
ForceNew: true,
ValidateFunc: validation.StringInSlice([]string{"S3", "S3gov", "GCS", "AZURE", "S3GOV"}, false),
Type: schema.TypeString,
Required: true,
ForceNew: true,
ValidateDiagFunc: StringInSlice(append(sdk.AsStringList(sdk.AllS3Protocols), "GCS", "AZURE"), true),
},
"storage_aws_external_id": {
Type: schema.TypeString,
Expand Down Expand Up @@ -140,7 +140,7 @@ func CreateStorageIntegration(d *schema.ResourceData, meta any) error {
req := sdk.NewCreateStorageIntegrationRequest(name, enabled, storageAllowedLocations)

if v, ok := d.GetOk("comment"); ok {
req.WithComment(sdk.String(v.(string)))
req.WithComment(v.(string))
}

if _, ok := d.GetOk("storage_blocked_locations"); ok {
Expand All @@ -154,28 +154,33 @@ func CreateStorageIntegration(d *schema.ResourceData, meta any) error {
req.WithStorageBlockedLocations(storageBlockedLocations)
}

storageProvider := d.Get("storage_provider").(string)
storageProvider := strings.ToUpper(d.Get("storage_provider").(string))

switch {
case slices.Contains(sdk.AllS3Protocols, sdk.S3Protocol(storageProvider)):
s3Protocol, err := sdk.ToS3Protocol(storageProvider)
if err != nil {
return err
}

switch storageProvider {
case "S3", "S3GOV", "S3gov":
v, ok := d.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")
}

s3Params := sdk.NewS3StorageParamsRequest(v.(string))
s3Params := sdk.NewS3StorageParamsRequest(s3Protocol, v.(string))
if _, ok := d.GetOk("storage_aws_object_acl"); ok {
s3Params.WithStorageAwsObjectAcl(sdk.String(d.Get("storage_aws_object_acl").(string)))
s3Params.WithStorageAwsObjectAcl(d.Get("storage_aws_object_acl").(string))
}
req.WithS3StorageProviderParams(s3Params)
case "AZURE":
req.WithS3StorageProviderParams(*s3Params)
case storageProvider == "AZURE":
v, ok := d.GetOk("azure_tenant_id")
if !ok {
return fmt.Errorf("if you use the Azure storage provider you must specify an azure_tenant_id")
}
req.WithAzureStorageProviderParams(sdk.NewAzureStorageParamsRequest(sdk.String(v.(string))))
case "GCS":
req.WithGCSStorageProviderParams(sdk.NewGCSStorageParamsRequest())
req.WithAzureStorageProviderParams(*sdk.NewAzureStorageParamsRequest(sdk.String(v.(string))))
case storageProvider == "GCS":
req.WithGCSStorageProviderParams(*sdk.NewGCSStorageParamsRequest())
default:
return fmt.Errorf("unexpected provider %v", storageProvider)
}
Expand Down Expand Up @@ -295,7 +300,7 @@ func UpdateStorageIntegration(d *schema.ResourceData, meta any) error {

if d.HasChange("comment") {
runSetStatement = true
setReq.WithComment(sdk.String(d.Get("comment").(string)))
setReq.WithComment(d.Get("comment").(string))
}

if d.HasChange("enabled") {
Expand All @@ -320,7 +325,7 @@ func UpdateStorageIntegration(d *schema.ResourceData, meta any) error {
v := d.Get("storage_blocked_locations").([]interface{})
if len(v) == 0 {
if err := client.StorageIntegrations.Alter(ctx, sdk.NewAlterStorageIntegrationRequest(id).
WithUnset(sdk.NewStorageIntegrationUnsetRequest().WithStorageBlockedLocations(sdk.Bool(true)))); err != nil {
WithUnset(*sdk.NewStorageIntegrationUnsetRequest().WithStorageBlockedLocations(true))); err != nil {
return fmt.Errorf("error unsetting storage_blocked_locations, err = %w", err)
}
} else {
Expand All @@ -342,25 +347,25 @@ func UpdateStorageIntegration(d *schema.ResourceData, meta any) error {

if d.HasChange("storage_aws_object_acl") {
if v, ok := d.GetOk("storage_aws_object_acl"); ok {
s3SetParams.WithStorageAwsObjectAcl(sdk.String(v.(string)))
s3SetParams.WithStorageAwsObjectAcl(v.(string))
} else {
if err := client.StorageIntegrations.Alter(ctx, sdk.NewAlterStorageIntegrationRequest(id).
WithUnset(sdk.NewStorageIntegrationUnsetRequest().WithStorageAwsObjectAcl(sdk.Bool(true)))); err != nil {
WithUnset(*sdk.NewStorageIntegrationUnsetRequest().WithStorageAwsObjectAcl(true))); err != nil {
return fmt.Errorf("error unsetting storage_aws_object_acl, err = %w", err)
}
}
}

setReq.WithS3Params(s3SetParams)
setReq.WithS3Params(*s3SetParams)
}

if d.HasChange("azure_tenant_id") {
runSetStatement = true
setReq.WithAzureParams(sdk.NewSetAzureStorageParamsRequest(d.Get("azure_tenant_id").(string)))
setReq.WithAzureParams(*sdk.NewSetAzureStorageParamsRequest(d.Get("azure_tenant_id").(string)))
}

if runSetStatement {
if err := client.StorageIntegrations.Alter(ctx, sdk.NewAlterStorageIntegrationRequest(id).WithSet(setReq)); err != nil {
if err := client.StorageIntegrations.Alter(ctx, sdk.NewAlterStorageIntegrationRequest(id).WithSet(*setReq)); err != nil {
return fmt.Errorf("error updating storage integration, err = %w", err)
}
}
Expand Down
2 changes: 1 addition & 1 deletion pkg/sdk/storage_integration_def.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ var StorageIntegrationDef = g.NewInterface(
TextAssignment("AZURE_TENANT_ID", g.ParameterOptions().SingleQuotes().Required()),
g.KeywordOptions(),
).
BooleanAssignment("ENABLED", g.ParameterOptions()).
OptionalBooleanAssignment("ENABLED", g.ParameterOptions()).
ListAssignment("STORAGE_ALLOWED_LOCATIONS", "StorageLocation", g.ParameterOptions().Parentheses()).
ListAssignment("STORAGE_BLOCKED_LOCATIONS", "StorageLocation", g.ParameterOptions().Parentheses()).
OptionalComment(),
Expand Down
82 changes: 42 additions & 40 deletions pkg/sdk/storage_integration_dto_builders_gen.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 2 additions & 1 deletion pkg/sdk/storage_integration_dto_gen.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,8 @@ type CreateStorageIntegrationRequest struct {
}

type S3StorageParamsRequest struct {
StorageAwsRoleArn string // required
Protocol S3Protocol // required
StorageAwsRoleArn string // required
StorageAwsObjectAcl *string
}

Expand Down
6 changes: 3 additions & 3 deletions pkg/sdk/storage_integration_gen.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,9 @@ type StorageLocation struct {
}

type S3StorageParams struct {
storageProvider string `ddl:"static" sql:"STORAGE_PROVIDER = 'S3'"`
StorageAwsRoleArn string `ddl:"parameter,single_quotes" sql:"STORAGE_AWS_ROLE_ARN"`
StorageAwsObjectAcl *string `ddl:"parameter,single_quotes" sql:"STORAGE_AWS_OBJECT_ACL"`
Protocol S3Protocol `ddl:"parameter,single_quotes" sql:"STORAGE_PROVIDER"`
StorageAwsRoleArn string `ddl:"parameter,single_quotes" sql:"STORAGE_AWS_ROLE_ARN"`
StorageAwsObjectAcl *string `ddl:"parameter,single_quotes" sql:"STORAGE_AWS_OBJECT_ACL"`
}

type GCSStorageParams struct {
Expand Down
14 changes: 14 additions & 0 deletions pkg/sdk/storage_integration_gen_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ func TestStorageIntegrations_Create(t *testing.T) {
return &CreateStorageIntegrationOptions{
name: id,
S3StorageProviderParams: &S3StorageParams{
Protocol: RegularS3Protocol,
StorageAwsRoleArn: "arn:aws:iam::001234567890:role/role",
},
Enabled: true,
Expand Down Expand Up @@ -56,10 +57,23 @@ func TestStorageIntegrations_Create(t *testing.T) {
assertOptsValidAndSQLEquals(t, opts, `CREATE STORAGE INTEGRATION %s TYPE = EXTERNAL_STAGE STORAGE_PROVIDER = 'S3' STORAGE_AWS_ROLE_ARN = 'arn:aws:iam::001234567890:role/role' ENABLED = true STORAGE_ALLOWED_LOCATIONS = ('allowed-loc-1', 'allowed-loc-2')`, id.FullyQualifiedName())
})

t.Run("basic - s3 gov protocol", func(t *testing.T) {
opts := defaultOpts()
opts.S3StorageProviderParams.Protocol = GovS3Protocol
assertOptsValidAndSQLEquals(t, opts, `CREATE STORAGE INTEGRATION %s TYPE = EXTERNAL_STAGE STORAGE_PROVIDER = 'S3GOV' STORAGE_AWS_ROLE_ARN = 'arn:aws:iam::001234567890:role/role' ENABLED = true STORAGE_ALLOWED_LOCATIONS = ('allowed-loc-1', 'allowed-loc-2')`, id.FullyQualifiedName())
})

t.Run("basic - s3 china protocol", func(t *testing.T) {
opts := defaultOpts()
opts.S3StorageProviderParams.Protocol = ChinaS3Protocol
assertOptsValidAndSQLEquals(t, opts, `CREATE STORAGE INTEGRATION %s TYPE = EXTERNAL_STAGE STORAGE_PROVIDER = 'S3CHINA' STORAGE_AWS_ROLE_ARN = 'arn:aws:iam::001234567890:role/role' ENABLED = true STORAGE_ALLOWED_LOCATIONS = ('allowed-loc-1', 'allowed-loc-2')`, id.FullyQualifiedName())
})

t.Run("all options - s3", func(t *testing.T) {
opts := defaultOpts()
opts.IfNotExists = Bool(true)
opts.S3StorageProviderParams = &S3StorageParams{
Protocol: RegularS3Protocol,
StorageAwsRoleArn: "arn:aws:iam::001234567890:role/role",
StorageAwsObjectAcl: String("bucket-owner-full-control"),
}
Expand Down
Loading

0 comments on commit c35bf68

Please sign in to comment.