From 66d3d1056f553368321c6c5d8cab496211129dab Mon Sep 17 00:00:00 2001 From: Ashen Gunaratne Date: Thu, 31 Oct 2024 15:41:07 +0000 Subject: [PATCH] Add support partitions in policy data sources --- aws/constants.go | 17 +++++ aws/data_aws_assume_role_policy.go | 26 ++++++-- aws/data_aws_assume_role_policy_test.go | 49 ++++++++++++++ aws/data_aws_bucket_policy.go | 25 +++++-- aws/data_aws_bucket_policy_test.go | 16 +++++ aws/data_aws_crossaccount_policy.go | 51 ++++++++------- aws/data_aws_crossaccount_policy_test.go | 29 +++++++++ ...ta_aws_unity_catalog_assume_role_policy.go | 14 +++- ...s_unity_catalog_assume_role_policy_test.go | 65 +++++++++++++++++++ aws/data_aws_unity_catalog_policy.go | 17 +++-- aws/data_aws_unity_catalog_policy_test.go | 58 +++++++++++++++++ docs/data-sources/aws_assume_role_policy.md | 1 + docs/data-sources/aws_bucket_policy.md | 1 + docs/data-sources/aws_crossaccount_policy.md | 1 + .../aws_unity_catalog_assume_role_policy.md | 3 +- docs/data-sources/aws_unity_catalog_policy.md | 1 + 16 files changed, 330 insertions(+), 44 deletions(-) create mode 100644 aws/constants.go diff --git a/aws/constants.go b/aws/constants.go new file mode 100644 index 0000000000..36d9f84ea1 --- /dev/null +++ b/aws/constants.go @@ -0,0 +1,17 @@ +package aws + +var AwsConfig = map[string]map[string]string{ + "aws": { + "accountId": "414351767826", + "logDeliveryIamArn": "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK", + "unityCatalogueIamArn": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL", + }, + "aws-us-gov": { + "accountId": "044793339203", + "logDeliveryIamArn": "arn:aws-us-gov:iam::044793339203:role/SaasUsageDeliveryRole-prod-aws-gov-IAMRole-L4QM0RCHYQ1G", + "unityCatalogueIamArn": "arn:aws-us-gov:iam::044793339203:role/unity-catalog-prod-UCMasterRole-1QRFA8SGY15OJ", + }, +} + +var AwsPartitions = []string{"aws", "aws-us-gov"} +var AwsPartitionsValidationError = "aws_partition must be either 'aws' or 'aws-us-gov'" diff --git a/aws/data_aws_assume_role_policy.go b/aws/data_aws_assume_role_policy.go index 576321d819..1cbbed669d 100644 --- a/aws/data_aws_assume_role_policy.go +++ b/aws/data_aws_assume_role_policy.go @@ -7,6 +7,7 @@ import ( "github.com/databricks/terraform-provider-databricks/common" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation" ) type awsIamPolicy struct { @@ -31,6 +32,13 @@ func DataAwsAssumeRolePolicy() common.Resource { return common.Resource{ Read: func(ctx context.Context, d *schema.ResourceData, m *common.DatabricksClient) error { externalID := d.Get("external_id").(string) + awsPartition := d.Get("aws_partition").(string) + databricksAwsAccountId := d.Get("databricks_account_id").(string) + + if databricksAwsAccountId == "" { + databricksAwsAccountId = AwsConfig[awsPartition]["accountId"] + } + policy := awsIamPolicy{ Version: "2012-10-17", Statements: []*awsIamPolicyStatement{ @@ -43,16 +51,14 @@ func DataAwsAssumeRolePolicy() common.Resource { }, }, Principal: map[string]string{ - "AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)), + "AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId), }, }, }, } if v, ok := d.GetOk("for_log_delivery"); ok { if v.(bool) { - // this is production UsageDelivery IAM role, that is considered a constant - logDeliveryARN := "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK" - policy.Statements[0].Principal["AWS"] = logDeliveryARN + policy.Statements[0].Principal["AWS"] = AwsConfig[awsPartition]["logDeliveryIamArn"] } } policyJSON, err := json.MarshalIndent(policy, "", " ") @@ -65,10 +71,16 @@ func DataAwsAssumeRolePolicy() common.Resource { return nil }, Schema: map[string]*schema.Schema{ + "aws_partition": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringInSlice(AwsPartitions, false), + Default: "aws", + }, "databricks_account_id": { - Type: schema.TypeString, - Default: "414351767826", - Optional: true, + Type: schema.TypeString, + Optional: true, + Deprecated: "databricks_account_id will be will be removed in the next major release.", }, "for_log_delivery": { Type: schema.TypeBool, diff --git a/aws/data_aws_assume_role_policy_test.go b/aws/data_aws_assume_role_policy_test.go index 7322660420..f4a1fa1998 100644 --- a/aws/data_aws_assume_role_policy_test.go +++ b/aws/data_aws_assume_role_policy_test.go @@ -19,3 +19,52 @@ func TestDataAwsAssumeRolePolicy(t *testing.T) { j := d.Get("json") assert.Lenf(t, j, 299, "Strange length for policy: %s", j) } + +func TestDataAwsAssumeRolePolicyGov(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsAssumeRolePolicy(), + NonWritable: true, + ID: ".", + HCL: ` + aws_partition = "aws-us-gov" + external_id = "abc" + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json") + assert.Lenf(t, j, 306, "Strange length for policy: %s", j) +} + +func TestDataAwsAssumeRolePolicyLogDelivery(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsAssumeRolePolicy(), + NonWritable: true, + ID: ".", + HCL: ` + external_id = "abc" + for_log_delivery = true + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json") + assert.Lenf(t, j, 347, "Strange length for policy: %s", j) +} + +func TestDataAwsAssumeRolePolicyLogDeliveryGov(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsAssumeRolePolicy(), + NonWritable: true, + ID: ".", + HCL: ` + aws_partition = "aws-us-gov" + external_id = "abc" + for_log_delivery = true + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json") + assert.Lenf(t, j, 362, "Strange length for policy: %s", j) +} diff --git a/aws/data_aws_bucket_policy.go b/aws/data_aws_bucket_policy.go index b1cc42a8f4..dc8394d85e 100644 --- a/aws/data_aws_bucket_policy.go +++ b/aws/data_aws_bucket_policy.go @@ -16,6 +16,13 @@ func DataAwsBucketPolicy() common.Resource { return common.Resource{ Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error { bucket := d.Get("bucket").(string) + awsPartition := d.Get("aws_partition").(string) + databricksAwsAccountId := AwsConfig[awsPartition]["accountId"] + + if databricksAwsAccountId == "" { + databricksAwsAccountId = AwsConfig[awsPartition]["accountId"] + } + policy := awsIamPolicy{ Version: "2012-10-17", Statements: []*awsIamPolicyStatement{ @@ -30,11 +37,11 @@ func DataAwsBucketPolicy() common.Resource { "s3:GetBucketLocation", }, Resources: []string{ - fmt.Sprintf("arn:aws:s3:::%s/*", bucket), - fmt.Sprintf("arn:aws:s3:::%s", bucket), + fmt.Sprintf("arn:%s:s3:::%s/*", awsPartition, bucket), + fmt.Sprintf("arn:%s:s3:::%s", awsPartition, bucket), }, Principal: map[string]string{ - "AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)), + "AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId), }, }, }, @@ -60,10 +67,16 @@ func DataAwsBucketPolicy() common.Resource { return nil }, Schema: map[string]*schema.Schema{ + "aws_partition": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringInSlice(AwsPartitions, false), + Default: "aws", + }, "databricks_account_id": { - Type: schema.TypeString, - Default: "414351767826", - Optional: true, + Type: schema.TypeString, + Optional: true, + Deprecated: "databricks_account_id will be will be removed in the next major release.", }, "databricks_e2_account_id": { Type: schema.TypeString, diff --git a/aws/data_aws_bucket_policy_test.go b/aws/data_aws_bucket_policy_test.go index 75f3a13645..5ec6c763b2 100644 --- a/aws/data_aws_bucket_policy_test.go +++ b/aws/data_aws_bucket_policy_test.go @@ -53,3 +53,19 @@ func TestDataAwsBucketPolicyConfusedDeputyProblem(t *testing.T) { j := d.Get("json") assert.Lenf(t, j, 575, "Strange length for policy: %s", j) } + +func TestDataAwsBucketPolicyPartitionGov(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsBucketPolicy(), + NonWritable: true, + ID: ".", + HCL: ` + bucket = "abc" + aws_partition = "aws-us-gov" + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json") + assert.Lenf(t, j, 461, "Strange length for policy: %s", j) +} diff --git a/aws/data_aws_crossaccount_policy.go b/aws/data_aws_crossaccount_policy.go index a5da5d9365..46ff5a6b9e 100644 --- a/aws/data_aws_crossaccount_policy.go +++ b/aws/data_aws_crossaccount_policy.go @@ -3,6 +3,7 @@ package aws import ( "context" "encoding/json" + "errors" "fmt" "regexp" "slices" @@ -17,11 +18,16 @@ func DataAwsCrossaccountPolicy() common.Resource { PassRole []string `json:"pass_roles,omitempty"` JSON string `json:"json" tf:"computed"` AwsAccountId string `json:"aws_account_id,omitempty"` + AwsPartition string `json:"aws_partition,omitempty" tf:"default:aws"` VpcId string `json:"vpc_id,omitempty"` Region string `json:"region,omitempty"` SecurityGroupId string `json:"security_group_id,omitempty"` } return common.NoClientData(func(ctx context.Context, data *AwsCrossAccountPolicy) error { + if !slices.Contains(AwsPartitions, data.AwsPartition) { + return errors.New(AwsPartitionsValidationError) + } + if !slices.Contains([]string{"managed", "customer", "restricted"}, data.PolicyType) { return fmt.Errorf("policy_type must be either 'managed', 'customer' or 'restricted'") } @@ -145,7 +151,7 @@ func DataAwsCrossaccountPolicy() common.Resource { "iam:CreateServiceLinkedRole", "iam:PutRolePolicy", }, - Resources: "arn:aws:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot", + Resources: fmt.Sprintf("arn:%s:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot", data.AwsPartition), Condition: map[string]map[string]string{ "StringLike": { "iam:AWSServiceName": "spot.amazonaws.com", @@ -168,6 +174,7 @@ func DataAwsCrossaccountPolicy() common.Resource { if data.PolicyType == "restricted" { region := data.Region aws_account_id := data.AwsAccountId + awsPartition := data.AwsPartition vpc_id := data.VpcId security_group_id := data.SecurityGroupId policy.Statements = append(policy.Statements, @@ -179,7 +186,7 @@ func DataAwsCrossaccountPolicy() common.Resource { "ec2:DisassociateIamInstanceProfile", "ec2:ReplaceIamInstanceProfileAssociation", }, - Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id), + Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id), Condition: map[string]map[string]string{ "StringEquals": { "ec2:ResourceTag/Vendor": "Databricks", @@ -191,8 +198,8 @@ func DataAwsCrossaccountPolicy() common.Resource { Effect: "Allow", Actions: "ec2:RunInstances", Resources: []string{ - fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id), }, Condition: map[string]map[string]string{ "StringEquals": { @@ -204,7 +211,7 @@ func DataAwsCrossaccountPolicy() common.Resource { Sid: "AllowEc2RunInstanceImagePerTag", Effect: "Allow", Actions: "ec2:RunInstances", - Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id), + Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id), Condition: map[string]map[string]string{ "StringEquals": { "aws:ResourceTag/Vendor": "Databricks", @@ -216,13 +223,13 @@ func DataAwsCrossaccountPolicy() common.Resource { Effect: "Allow", Actions: "ec2:RunInstances", Resources: []string{ - fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id), }, Condition: map[string]map[string]string{ "StringEquals": { - "ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id), + "ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id), }, }, }, @@ -231,19 +238,19 @@ func DataAwsCrossaccountPolicy() common.Resource { Effect: "Allow", Actions: "ec2:RunInstances", NotResources: []string{ - fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id), }, }, &awsIamPolicyStatement{ Sid: "EC2TerminateInstancesTag", Effect: "Allow", Actions: "ec2:TerminateInstances", - Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id), + Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id), Condition: map[string]map[string]string{ "StringEquals": { "ec2:ResourceTag/Vendor": "Databricks", @@ -258,8 +265,8 @@ func DataAwsCrossaccountPolicy() common.Resource { "ec2:DetachVolume", }, Resources: []string{ - fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id), - fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id), }, Condition: map[string]map[string]string{ "StringEquals": { @@ -271,7 +278,7 @@ func DataAwsCrossaccountPolicy() common.Resource { Sid: "EC2CreateVolumeByTag", Effect: "Allow", Actions: "ec2:CreateVolume", - Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id), + Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id), Condition: map[string]map[string]string{ "StringEquals": { "aws:RequestTag/Vendor": "Databricks", @@ -283,7 +290,7 @@ func DataAwsCrossaccountPolicy() common.Resource { Effect: "Allow", Actions: "ec2:DeleteVolume", Resources: []string{ - fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id), + fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id), }, Condition: map[string]map[string]string{ "StringEquals": { @@ -300,10 +307,10 @@ func DataAwsCrossaccountPolicy() common.Resource { "ec2:RevokeSecurityGroupEgress", "ec2:RevokeSecurityGroupIngress", }, - Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/%s", region, aws_account_id, security_group_id), + Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/%s", awsPartition, region, aws_account_id, security_group_id), Condition: map[string]map[string]string{ "StringEquals": { - "ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id), + "ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id), }, }, }, diff --git a/aws/data_aws_crossaccount_policy_test.go b/aws/data_aws_crossaccount_policy_test.go index 177cb166e9..6832807be5 100644 --- a/aws/data_aws_crossaccount_policy_test.go +++ b/aws/data_aws_crossaccount_policy_test.go @@ -530,6 +530,25 @@ func TestDataAwsCrossAccountRestrictedPolicy(t *testing.T) { assert.Lenf(t, j, 5725, "Strange length for policy: %s", j) } +func TestDataAwsCrossAccountRestrictedPolicyPartitionGov(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsCrossaccountPolicy(), + NonWritable: true, + HCL: ` + policy_type = "restricted" + aws_account_id = "123456789012" + aws_partition = "aws-us-gov" + vpc_id = "vpc-12345678" + region = "us-west-2" + security_group_id = "sg-12345678"`, + ID: ".", + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json") + assert.Lenf(t, j, 5879, "Strange length for policy: %s", j) +} + func TestDataAwsCrossAccountInvalidPolicy(t *testing.T) { qa.ResourceFixture{ Read: true, @@ -552,6 +571,16 @@ func TestDataAwsCrossAccountInvalidAccountId(t *testing.T) { }.ExpectError(t, "aws_account_id must be a 12 digit number") } +func TestDataAwsCrossAccountInvalidPartition(t *testing.T) { + qa.ResourceFixture{ + Read: true, + Resource: DataAwsCrossaccountPolicy(), + NonWritable: true, + HCL: `aws_partition = "something"`, + ID: ".", + }.ExpectError(t, AwsPartitionsValidationError) +} + func TestDataAwsCrossAccountInvalidVpcId(t *testing.T) { qa.ResourceFixture{ Read: true, diff --git a/aws/data_aws_unity_catalog_assume_role_policy.go b/aws/data_aws_unity_catalog_assume_role_policy.go index d4706bdca5..a90ab98505 100644 --- a/aws/data_aws_unity_catalog_assume_role_policy.go +++ b/aws/data_aws_unity_catalog_assume_role_policy.go @@ -3,7 +3,9 @@ package aws import ( "context" "encoding/json" + "errors" "fmt" + "slices" "github.com/databricks/terraform-provider-databricks/common" ) @@ -14,13 +16,19 @@ func DataAwsUnityCatalogAssumeRolePolicy() common.Resource { UnityCatalogIamArn string `json:"unity_catalog_iam_arn,omitempty" tf:"computed"` ExternalId string `json:"external_id"` AwsAccountId string `json:"aws_account_id"` + AwsPartition string `json:"aws_partition,omitempty" tf:"default:aws"` JSON string `json:"json" tf:"computed"` Id string `json:"id" tf:"computed"` } return common.NoClientData(func(ctx context.Context, data *AwsUcAssumeRolePolicy) error { + if !slices.Contains(AwsPartitions, data.AwsPartition) { + return errors.New(AwsPartitionsValidationError) + } + if data.UnityCatalogIamArn == "" { - data.UnityCatalogIamArn = "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL" + data.UnityCatalogIamArn = AwsConfig[data.AwsPartition]["unityCatalogueIamArn"] } + policy := awsIamPolicy{ Version: "2012-10-17", Statements: []*awsIamPolicyStatement{ @@ -43,11 +51,11 @@ func DataAwsUnityCatalogAssumeRolePolicy() common.Resource { Actions: "sts:AssumeRole", Condition: map[string]map[string]string{ "ArnLike": { - "aws:PrincipalArn": fmt.Sprintf("arn:aws:iam::%s:role/%s", data.AwsAccountId, data.RoleName), + "aws:PrincipalArn": fmt.Sprintf("arn:%s:iam::%s:role/%s", data.AwsPartition, data.AwsAccountId, data.RoleName), }, }, Principal: map[string]string{ - "AWS": fmt.Sprintf("arn:aws:iam::%s:root", data.AwsAccountId), + "AWS": fmt.Sprintf("arn:%s:iam::%s:root", data.AwsPartition, data.AwsAccountId), }, }, }, diff --git a/aws/data_aws_unity_catalog_assume_role_policy_test.go b/aws/data_aws_unity_catalog_assume_role_policy_test.go index 30c1d89f2d..4f2da8932e 100644 --- a/aws/data_aws_unity_catalog_assume_role_policy_test.go +++ b/aws/data_aws_unity_catalog_assume_role_policy_test.go @@ -103,3 +103,68 @@ func TestDataAwsUnityCatalogAssumeRolePolicyWithoutUcArn(t *testing.T) { }` compareJSON(t, j, p) } + +func TestDataAwsUnityCatalogAssumeRolePolicyGovWithoutUcArn(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsUnityCatalogAssumeRolePolicy(), + NonWritable: true, + ID: ".", + HCL: ` + aws_account_id = "123456789098" + aws_partition = "aws-us-gov" + role_name = "databricks-role" + external_id = "12345" + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json").(string) + p := `{ + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "UnityCatalogAssumeRole", + "Effect": "Allow", + "Action": "sts:AssumeRole", + "Principal": { + "AWS": "arn:aws-us-gov:iam::044793339203:role/unity-catalog-prod-UCMasterRole-1QRFA8SGY15OJ" + }, + "Condition": { + "StringEquals": { + "sts:ExternalId": "12345" + } + } + }, + { + "Sid": "ExplicitSelfRoleAssumption", + "Effect": "Allow", + "Action": "sts:AssumeRole", + "Principal": { + "AWS": "arn:aws-us-gov:iam::123456789098:root" + }, + "Condition": { + "ArnLike": { + "aws:PrincipalArn": "arn:aws-us-gov:iam::123456789098:role/databricks-role" + } + } + } + ] + }` + compareJSON(t, j, p) +} + +func TestDataAwsUnityCatalogAssumeRolePolicyInvalidPartition(t *testing.T) { + qa.ResourceFixture{ + Read: true, + Resource: DataAwsUnityCatalogAssumeRolePolicy(), + NonWritable: true, + ID: ".", + HCL: ` + aws_account_id = "123456789098" + aws_partition = "something" + role_name = "databricks-role" + unity_catalog_iam_arn = "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL" + external_id = "12345" + `, + }.ExpectError(t, AwsPartitionsValidationError) +} diff --git a/aws/data_aws_unity_catalog_policy.go b/aws/data_aws_unity_catalog_policy.go index d332b84348..5dbc565b38 100644 --- a/aws/data_aws_unity_catalog_policy.go +++ b/aws/data_aws_unity_catalog_policy.go @@ -15,6 +15,7 @@ import ( func generateReadContext(ctx context.Context, d *schema.ResourceData, m *common.DatabricksClient) error { bucket := d.Get("bucket_name").(string) awsAccountId := d.Get("aws_account_id").(string) + awsPartition := d.Get("aws_partition").(string) roleName := d.Get("role_name").(string) policy := awsIamPolicy{ Version: "2012-10-17", @@ -29,8 +30,8 @@ func generateReadContext(ctx context.Context, d *schema.ResourceData, m *common. "s3:GetBucketLocation", }, Resources: []string{ - fmt.Sprintf("arn:aws:s3:::%s/*", bucket), - fmt.Sprintf("arn:aws:s3:::%s", bucket), + fmt.Sprintf("arn:%s:s3:::%s/*", awsPartition, bucket), + fmt.Sprintf("arn:%s:s3:::%s", awsPartition, bucket), }, }, { @@ -39,14 +40,14 @@ func generateReadContext(ctx context.Context, d *schema.ResourceData, m *common. "sts:AssumeRole", }, Resources: []string{ - fmt.Sprintf("arn:aws:iam::%s:role/%s", awsAccountId, roleName), + fmt.Sprintf("arn:%s:iam::%s:role/%s", awsPartition, awsAccountId, roleName), }, }, }, } if kmsKey, ok := d.GetOk("kms_name"); ok { - kmsArn := fmt.Sprintf("arn:aws:kms:%s", kmsKey) - if strings.HasPrefix(kmsKey.(string), "arn:aws") { + kmsArn := fmt.Sprintf("arn:%s:kms:%s", awsPartition, kmsKey) + if strings.HasPrefix(kmsKey.(string), fmt.Sprintf("arn:%s", awsPartition)) { kmsArn = kmsKey.(string) } policy.Statements = append(policy.Statements, &awsIamPolicyStatement{ @@ -92,6 +93,12 @@ func validateSchema() map[string]*schema.Schema { Type: schema.TypeString, Required: true, }, + "aws_partition": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringInSlice(AwsPartitions, false), + Default: "aws", + }, "json": { Type: schema.TypeString, Computed: true, diff --git a/aws/data_aws_unity_catalog_policy_test.go b/aws/data_aws_unity_catalog_policy_test.go index 28a45a4f16..6ca159e290 100644 --- a/aws/data_aws_unity_catalog_policy_test.go +++ b/aws/data_aws_unity_catalog_policy_test.go @@ -167,6 +167,64 @@ func TestDataAwsUnityCatalogPolicyWithoutKMS(t *testing.T) { compareJSON(t, j, p) } +func TestDataAwsUnityCatalogPolicyPartionGov(t *testing.T) { + d, err := qa.ResourceFixture{ + Read: true, + Resource: DataAwsUnityCatalogPolicy(), + NonWritable: true, + ID: ".", + HCL: ` + aws_account_id = "123456789098" + aws_partition = "aws-us-gov" + bucket_name = "databricks-bucket" + role_name = "databricks-role" + kms_name = "databricks-kms" + `, + }.Apply(t) + assert.NoError(t, err) + j := d.Get("json").(string) + p := `{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "s3:GetObject", + "s3:PutObject", + "s3:DeleteObject", + "s3:ListBucket", + "s3:GetBucketLocation" + ], + "Resource": [ + "arn:aws-us-gov:s3:::databricks-bucket/*", + "arn:aws-us-gov:s3:::databricks-bucket" + ] + }, + { + "Effect": "Allow", + "Action": [ + "sts:AssumeRole" + ], + "Resource": [ + "arn:aws-us-gov:iam::123456789098:role/databricks-role" + ] + }, + { + "Effect": "Allow", + "Action": [ + "kms:Decrypt", + "kms:Encrypt", + "kms:GenerateDataKey*" + ], + "Resource": [ + "arn:aws-us-gov:kms:databricks-kms" + ] + } + ] + }` + compareJSON(t, j, p) +} + func compareJSON(t *testing.T, json1 string, json2 string) { var i1 interface{} var i2 interface{} diff --git a/docs/data-sources/aws_assume_role_policy.md b/docs/data-sources/aws_assume_role_policy.md index 73d6fb0e11..d46f3520b7 100644 --- a/docs/data-sources/aws_assume_role_policy.md +++ b/docs/data-sources/aws_assume_role_policy.md @@ -49,6 +49,7 @@ resource "databricks_mws_credentials" "this" { ## Argument Reference * `external_id` (Required) Account Id that could be found in the top right corner of [Accounts Console](https://accounts.cloud.databricks.com/). +* `aws_partition` - (Optional) AWS partition. The options are `aws` or `aws-us-gov`. Defaults to `aws` * `for_log_delivery` (Optional) Either or not this assume role policy should be created for usage log delivery. Defaults to false. ## Attribute Reference diff --git a/docs/data-sources/aws_bucket_policy.md b/docs/data-sources/aws_bucket_policy.md index e42949e06a..f2e99edc2a 100644 --- a/docs/data-sources/aws_bucket_policy.md +++ b/docs/data-sources/aws_bucket_policy.md @@ -75,6 +75,7 @@ resource "aws_s3_bucket_policy" "ds" { ## Argument Reference * `bucket` - (Required) AWS S3 Bucket name for which to generate the policy document. +* `aws_partition` - (Optional) AWS partition. The options are `aws` or `aws-us-gov`. Defaults to `aws` * `full_access_role` - (Optional) Data access role that can have full access for this bucket * `databricks_e2_account_id` - (Optional) Your Databricks account ID. Used to generate restrictive IAM policies that will increase the security of your root bucket diff --git a/docs/data-sources/aws_crossaccount_policy.md b/docs/data-sources/aws_crossaccount_policy.md index 715cf59b15..883bd1b490 100644 --- a/docs/data-sources/aws_crossaccount_policy.md +++ b/docs/data-sources/aws_crossaccount_policy.md @@ -21,6 +21,7 @@ data "databricks_aws_crossaccount_policy" "this" {} * `pass_roles` (Optional) (List) List of Data IAM role ARNs that are explicitly granted `iam:PassRole` action. The below arguments are only valid for `restricted` policy type * `aws_account_id` — Your AWS account ID, which is a number. +* `aws_partition` - (Optional) AWS partition. The options are `aws` or `aws-us-gov`. Defaults to `aws` * `vpc_id` — ID of the AWS VPC where you want to launch workspaces. * `region` — AWS Region name for your VPC deployment, for example `us-west-2`. * `security_group_id` — ID of your AWS security group. When you add a security group restriction, you cannot reuse the cross-account IAM role or reference a credentials ID (`credentials_id`) for any other workspaces. For those other workspaces, you must create separate roles, policies, and credentials objects. diff --git a/docs/data-sources/aws_unity_catalog_assume_role_policy.md b/docs/data-sources/aws_unity_catalog_assume_role_policy.md index c5f66ddec9..1619855ca9 100644 --- a/docs/data-sources/aws_unity_catalog_assume_role_policy.md +++ b/docs/data-sources/aws_unity_catalog_assume_role_policy.md @@ -38,9 +38,10 @@ resource "aws_iam_role" "metastore_data_access" { ## Argument Reference * `aws_account_id` (Required) The Account ID of the current AWS account (not your Databricks account). +* `aws_partition` - (Optional) AWS partition. The options are `aws` or `aws-us-gov`. Defaults to `aws` * `external_id` (Required) The [storage credential](../resources/storage_credential.md) external id. * `role_name` (Required) The name of the AWS IAM role to be created for Unity Catalog. -* `unity_catalog_iam_arn` (Optional) The Databricks Unity Catalog IAM Role ARN. Defaults to `arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL` +* `unity_catalog_iam_arn` (Optional) The Databricks Unity Catalog IAM Role ARN. Defaults to `arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL` on standard AWS partition selection and `arn:aws-us-gov:iam::044793339203:role/unity-catalog-prod-UCMasterRole-1QRFA8SGY15OJ` on GovCloud partition selection ## Attribute Reference diff --git a/docs/data-sources/aws_unity_catalog_policy.md b/docs/data-sources/aws_unity_catalog_policy.md index 3804b1d5fa..2e65039d57 100644 --- a/docs/data-sources/aws_unity_catalog_policy.md +++ b/docs/data-sources/aws_unity_catalog_policy.md @@ -38,6 +38,7 @@ resource "aws_iam_role" "metastore_data_access" { ## Argument Reference * `aws_account_id` (Required) The Account ID of the current AWS account (not your Databricks account). +* `aws_partition` - (Optional) AWS partition. The options are `aws` or `aws-us-gov`. Defaults to `aws` * `bucket_name` (Required) The name of the S3 bucket used as root storage location for [managed tables](https://docs.databricks.com/data-governance/unity-catalog/index.html#managed-table) in Unity Catalog. * `role_name` (Required) The name of the AWS IAM role that you created in the previous step in the [official documentation](https://docs.databricks.com/data-governance/unity-catalog/get-started.html#configure-a-storage-bucket-and-iam-role-in-aws). * `kms_name` (Optional) If encryption is enabled, provide the ARN of the KMS key that encrypts the S3 bucket contents. If encryption is disabled, do not provide this argument.