Skip to content

Commit d4e461c

Browse files
ashenmashenm
authored
[Feature] Add support partitions in policy data sources (databricks#4181)
## Changes - Resolves databricks#4054 - Resolves databricks#4152 - Add optional argument `aws_partition` to all aws policy data sources to allow usage in all aws partitions ## Tests <!-- How is this tested? Please see the checklist below and also describe any other relevant tests --> - [x] `make test` run locally - [x] relevant change in `docs/` folder - [x] covered with integration tests in `internal/acceptance` - [x] relevant acceptance tests are passing - [ ] using Go SDK
1 parent ae65156 commit d4e461c

16 files changed

+330
-44
lines changed

‎aws/constants.go

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package aws
2+
3+
var AwsConfig = map[string]map[string]string{
4+
"aws": {
5+
"accountId": "414351767826",
6+
"logDeliveryIamArn": "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK",
7+
"unityCatalogueIamArn": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL",
8+
},
9+
"aws-us-gov": {
10+
"accountId": "044793339203",
11+
"logDeliveryIamArn": "arn:aws-us-gov:iam::044793339203:role/SaasUsageDeliveryRole-prod-aws-gov-IAMRole-L4QM0RCHYQ1G",
12+
"unityCatalogueIamArn": "arn:aws-us-gov:iam::044793339203:role/unity-catalog-prod-UCMasterRole-1QRFA8SGY15OJ",
13+
},
14+
}
15+
16+
var AwsPartitions = []string{"aws", "aws-us-gov"}
17+
var AwsPartitionsValidationError = "aws_partition must be either 'aws' or 'aws-us-gov'"

‎aws/data_aws_assume_role_policy.go

+19-7
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77

88
"github.com/databricks/terraform-provider-databricks/common"
99
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
10+
"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
1011
)
1112

1213
type awsIamPolicy struct {
@@ -31,6 +32,13 @@ func DataAwsAssumeRolePolicy() common.Resource {
3132
return common.Resource{
3233
Read: func(ctx context.Context, d *schema.ResourceData, m *common.DatabricksClient) error {
3334
externalID := d.Get("external_id").(string)
35+
awsPartition := d.Get("aws_partition").(string)
36+
databricksAwsAccountId := d.Get("databricks_account_id").(string)
37+
38+
if databricksAwsAccountId == "" {
39+
databricksAwsAccountId = AwsConfig[awsPartition]["accountId"]
40+
}
41+
3442
policy := awsIamPolicy{
3543
Version: "2012-10-17",
3644
Statements: []*awsIamPolicyStatement{
@@ -43,16 +51,14 @@ func DataAwsAssumeRolePolicy() common.Resource {
4351
},
4452
},
4553
Principal: map[string]string{
46-
"AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)),
54+
"AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId),
4755
},
4856
},
4957
},
5058
}
5159
if v, ok := d.GetOk("for_log_delivery"); ok {
5260
if v.(bool) {
53-
// this is production UsageDelivery IAM role, that is considered a constant
54-
logDeliveryARN := "arn:aws:iam::414351767826:role/SaasUsageDeliveryRole-prod-IAMRole-3PLHICCRR1TK"
55-
policy.Statements[0].Principal["AWS"] = logDeliveryARN
61+
policy.Statements[0].Principal["AWS"] = AwsConfig[awsPartition]["logDeliveryIamArn"]
5662
}
5763
}
5864
policyJSON, err := json.MarshalIndent(policy, "", " ")
@@ -65,10 +71,16 @@ func DataAwsAssumeRolePolicy() common.Resource {
6571
return nil
6672
},
6773
Schema: map[string]*schema.Schema{
74+
"aws_partition": {
75+
Type: schema.TypeString,
76+
Optional: true,
77+
ValidateFunc: validation.StringInSlice(AwsPartitions, false),
78+
Default: "aws",
79+
},
6880
"databricks_account_id": {
69-
Type: schema.TypeString,
70-
Default: "414351767826",
71-
Optional: true,
81+
Type: schema.TypeString,
82+
Optional: true,
83+
Deprecated: "databricks_account_id will be will be removed in the next major release.",
7284
},
7385
"for_log_delivery": {
7486
Type: schema.TypeBool,

‎aws/data_aws_assume_role_policy_test.go

+49
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,52 @@ func TestDataAwsAssumeRolePolicy(t *testing.T) {
1919
j := d.Get("json")
2020
assert.Lenf(t, j, 299, "Strange length for policy: %s", j)
2121
}
22+
23+
func TestDataAwsAssumeRolePolicyGov(t *testing.T) {
24+
d, err := qa.ResourceFixture{
25+
Read: true,
26+
Resource: DataAwsAssumeRolePolicy(),
27+
NonWritable: true,
28+
ID: ".",
29+
HCL: `
30+
aws_partition = "aws-us-gov"
31+
external_id = "abc"
32+
`,
33+
}.Apply(t)
34+
assert.NoError(t, err)
35+
j := d.Get("json")
36+
assert.Lenf(t, j, 306, "Strange length for policy: %s", j)
37+
}
38+
39+
func TestDataAwsAssumeRolePolicyLogDelivery(t *testing.T) {
40+
d, err := qa.ResourceFixture{
41+
Read: true,
42+
Resource: DataAwsAssumeRolePolicy(),
43+
NonWritable: true,
44+
ID: ".",
45+
HCL: `
46+
external_id = "abc"
47+
for_log_delivery = true
48+
`,
49+
}.Apply(t)
50+
assert.NoError(t, err)
51+
j := d.Get("json")
52+
assert.Lenf(t, j, 347, "Strange length for policy: %s", j)
53+
}
54+
55+
func TestDataAwsAssumeRolePolicyLogDeliveryGov(t *testing.T) {
56+
d, err := qa.ResourceFixture{
57+
Read: true,
58+
Resource: DataAwsAssumeRolePolicy(),
59+
NonWritable: true,
60+
ID: ".",
61+
HCL: `
62+
aws_partition = "aws-us-gov"
63+
external_id = "abc"
64+
for_log_delivery = true
65+
`,
66+
}.Apply(t)
67+
assert.NoError(t, err)
68+
j := d.Get("json")
69+
assert.Lenf(t, j, 362, "Strange length for policy: %s", j)
70+
}

‎aws/data_aws_bucket_policy.go

+19-6
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,13 @@ func DataAwsBucketPolicy() common.Resource {
1616
return common.Resource{
1717
Read: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
1818
bucket := d.Get("bucket").(string)
19+
awsPartition := d.Get("aws_partition").(string)
20+
databricksAwsAccountId := AwsConfig[awsPartition]["accountId"]
21+
22+
if databricksAwsAccountId == "" {
23+
databricksAwsAccountId = AwsConfig[awsPartition]["accountId"]
24+
}
25+
1926
policy := awsIamPolicy{
2027
Version: "2012-10-17",
2128
Statements: []*awsIamPolicyStatement{
@@ -30,11 +37,11 @@ func DataAwsBucketPolicy() common.Resource {
3037
"s3:GetBucketLocation",
3138
},
3239
Resources: []string{
33-
fmt.Sprintf("arn:aws:s3:::%s/*", bucket),
34-
fmt.Sprintf("arn:aws:s3:::%s", bucket),
40+
fmt.Sprintf("arn:%s:s3:::%s/*", awsPartition, bucket),
41+
fmt.Sprintf("arn:%s:s3:::%s", awsPartition, bucket),
3542
},
3643
Principal: map[string]string{
37-
"AWS": fmt.Sprintf("arn:aws:iam::%s:root", d.Get("databricks_account_id").(string)),
44+
"AWS": fmt.Sprintf("arn:%s:iam::%s:root", awsPartition, databricksAwsAccountId),
3845
},
3946
},
4047
},
@@ -60,10 +67,16 @@ func DataAwsBucketPolicy() common.Resource {
6067
return nil
6168
},
6269
Schema: map[string]*schema.Schema{
70+
"aws_partition": {
71+
Type: schema.TypeString,
72+
Optional: true,
73+
ValidateFunc: validation.StringInSlice(AwsPartitions, false),
74+
Default: "aws",
75+
},
6376
"databricks_account_id": {
64-
Type: schema.TypeString,
65-
Default: "414351767826",
66-
Optional: true,
77+
Type: schema.TypeString,
78+
Optional: true,
79+
Deprecated: "databricks_account_id will be will be removed in the next major release.",
6780
},
6881
"databricks_e2_account_id": {
6982
Type: schema.TypeString,

‎aws/data_aws_bucket_policy_test.go

+16
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,19 @@ func TestDataAwsBucketPolicyConfusedDeputyProblem(t *testing.T) {
5353
j := d.Get("json")
5454
assert.Lenf(t, j, 575, "Strange length for policy: %s", j)
5555
}
56+
57+
func TestDataAwsBucketPolicyPartitionGov(t *testing.T) {
58+
d, err := qa.ResourceFixture{
59+
Read: true,
60+
Resource: DataAwsBucketPolicy(),
61+
NonWritable: true,
62+
ID: ".",
63+
HCL: `
64+
bucket = "abc"
65+
aws_partition = "aws-us-gov"
66+
`,
67+
}.Apply(t)
68+
assert.NoError(t, err)
69+
j := d.Get("json")
70+
assert.Lenf(t, j, 461, "Strange length for policy: %s", j)
71+
}

‎aws/data_aws_crossaccount_policy.go

+29-22
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package aws
33
import (
44
"context"
55
"encoding/json"
6+
"errors"
67
"fmt"
78
"regexp"
89
"slices"
@@ -17,11 +18,16 @@ func DataAwsCrossaccountPolicy() common.Resource {
1718
PassRole []string `json:"pass_roles,omitempty"`
1819
JSON string `json:"json" tf:"computed"`
1920
AwsAccountId string `json:"aws_account_id,omitempty"`
21+
AwsPartition string `json:"aws_partition,omitempty" tf:"default:aws"`
2022
VpcId string `json:"vpc_id,omitempty"`
2123
Region string `json:"region,omitempty"`
2224
SecurityGroupId string `json:"security_group_id,omitempty"`
2325
}
2426
return common.NoClientData(func(ctx context.Context, data *AwsCrossAccountPolicy) error {
27+
if !slices.Contains(AwsPartitions, data.AwsPartition) {
28+
return errors.New(AwsPartitionsValidationError)
29+
}
30+
2531
if !slices.Contains([]string{"managed", "customer", "restricted"}, data.PolicyType) {
2632
return fmt.Errorf("policy_type must be either 'managed', 'customer' or 'restricted'")
2733
}
@@ -145,7 +151,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
145151
"iam:CreateServiceLinkedRole",
146152
"iam:PutRolePolicy",
147153
},
148-
Resources: "arn:aws:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot",
154+
Resources: fmt.Sprintf("arn:%s:iam::*:role/aws-service-role/spot.amazonaws.com/AWSServiceRoleForEC2Spot", data.AwsPartition),
149155
Condition: map[string]map[string]string{
150156
"StringLike": {
151157
"iam:AWSServiceName": "spot.amazonaws.com",
@@ -168,6 +174,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
168174
if data.PolicyType == "restricted" {
169175
region := data.Region
170176
aws_account_id := data.AwsAccountId
177+
awsPartition := data.AwsPartition
171178
vpc_id := data.VpcId
172179
security_group_id := data.SecurityGroupId
173180
policy.Statements = append(policy.Statements,
@@ -179,7 +186,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
179186
"ec2:DisassociateIamInstanceProfile",
180187
"ec2:ReplaceIamInstanceProfileAssociation",
181188
},
182-
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
189+
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
183190
Condition: map[string]map[string]string{
184191
"StringEquals": {
185192
"ec2:ResourceTag/Vendor": "Databricks",
@@ -191,8 +198,8 @@ func DataAwsCrossaccountPolicy() common.Resource {
191198
Effect: "Allow",
192199
Actions: "ec2:RunInstances",
193200
Resources: []string{
194-
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
195-
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
201+
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
202+
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
196203
},
197204
Condition: map[string]map[string]string{
198205
"StringEquals": {
@@ -204,7 +211,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
204211
Sid: "AllowEc2RunInstanceImagePerTag",
205212
Effect: "Allow",
206213
Actions: "ec2:RunInstances",
207-
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id),
214+
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id),
208215
Condition: map[string]map[string]string{
209216
"StringEquals": {
210217
"aws:ResourceTag/Vendor": "Databricks",
@@ -216,13 +223,13 @@ func DataAwsCrossaccountPolicy() common.Resource {
216223
Effect: "Allow",
217224
Actions: "ec2:RunInstances",
218225
Resources: []string{
219-
fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id),
220-
fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id),
221-
fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id),
226+
fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id),
227+
fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id),
228+
fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id),
222229
},
223230
Condition: map[string]map[string]string{
224231
"StringEquals": {
225-
"ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id),
232+
"ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id),
226233
},
227234
},
228235
},
@@ -231,19 +238,19 @@ func DataAwsCrossaccountPolicy() common.Resource {
231238
Effect: "Allow",
232239
Actions: "ec2:RunInstances",
233240
NotResources: []string{
234-
fmt.Sprintf("arn:aws:ec2:%s:%s:image/*", region, aws_account_id),
235-
fmt.Sprintf("arn:aws:ec2:%s:%s:network-interface/*", region, aws_account_id),
236-
fmt.Sprintf("arn:aws:ec2:%s:%s:subnet/*", region, aws_account_id),
237-
fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/*", region, aws_account_id),
238-
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
239-
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
241+
fmt.Sprintf("arn:%s:ec2:%s:%s:image/*", awsPartition, region, aws_account_id),
242+
fmt.Sprintf("arn:%s:ec2:%s:%s:network-interface/*", awsPartition, region, aws_account_id),
243+
fmt.Sprintf("arn:%s:ec2:%s:%s:subnet/*", awsPartition, region, aws_account_id),
244+
fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/*", awsPartition, region, aws_account_id),
245+
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
246+
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
240247
},
241248
},
242249
&awsIamPolicyStatement{
243250
Sid: "EC2TerminateInstancesTag",
244251
Effect: "Allow",
245252
Actions: "ec2:TerminateInstances",
246-
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
253+
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
247254
Condition: map[string]map[string]string{
248255
"StringEquals": {
249256
"ec2:ResourceTag/Vendor": "Databricks",
@@ -258,8 +265,8 @@ func DataAwsCrossaccountPolicy() common.Resource {
258265
"ec2:DetachVolume",
259266
},
260267
Resources: []string{
261-
fmt.Sprintf("arn:aws:ec2:%s:%s:instance/*", region, aws_account_id),
262-
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
268+
fmt.Sprintf("arn:%s:ec2:%s:%s:instance/*", awsPartition, region, aws_account_id),
269+
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
263270
},
264271
Condition: map[string]map[string]string{
265272
"StringEquals": {
@@ -271,7 +278,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
271278
Sid: "EC2CreateVolumeByTag",
272279
Effect: "Allow",
273280
Actions: "ec2:CreateVolume",
274-
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
281+
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
275282
Condition: map[string]map[string]string{
276283
"StringEquals": {
277284
"aws:RequestTag/Vendor": "Databricks",
@@ -283,7 +290,7 @@ func DataAwsCrossaccountPolicy() common.Resource {
283290
Effect: "Allow",
284291
Actions: "ec2:DeleteVolume",
285292
Resources: []string{
286-
fmt.Sprintf("arn:aws:ec2:%s:%s:volume/*", region, aws_account_id),
293+
fmt.Sprintf("arn:%s:ec2:%s:%s:volume/*", awsPartition, region, aws_account_id),
287294
},
288295
Condition: map[string]map[string]string{
289296
"StringEquals": {
@@ -300,10 +307,10 @@ func DataAwsCrossaccountPolicy() common.Resource {
300307
"ec2:RevokeSecurityGroupEgress",
301308
"ec2:RevokeSecurityGroupIngress",
302309
},
303-
Resources: fmt.Sprintf("arn:aws:ec2:%s:%s:security-group/%s", region, aws_account_id, security_group_id),
310+
Resources: fmt.Sprintf("arn:%s:ec2:%s:%s:security-group/%s", awsPartition, region, aws_account_id, security_group_id),
304311
Condition: map[string]map[string]string{
305312
"StringEquals": {
306-
"ec2:vpc": fmt.Sprintf("arn:aws:ec2:%s:%s:vpc/%s", region, aws_account_id, vpc_id),
313+
"ec2:vpc": fmt.Sprintf("arn:%s:ec2:%s:%s:vpc/%s", awsPartition, region, aws_account_id, vpc_id),
307314
},
308315
},
309316
},

0 commit comments

Comments
 (0)