Skip to content

Commit 1c543c4

Browse files
authored
AWS role issue with external locations pointing to the root of a storage account (#3510)
Closes #3505
1 parent 8b8d66a commit 1c543c4

File tree

3 files changed

+31
-8
lines changed

3 files changed

+31
-8
lines changed

src/databricks/labs/ucx/assessment/aws.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ class AWSResources:
8787
S3_ACTIONS: typing.ClassVar[set[str]] = {"s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:PutObjectAcl"}
8888
S3_READONLY: typing.ClassVar[str] = "s3:GetObject"
8989
S3_REGEX: typing.ClassVar[str] = r"arn:aws:s3:::([a-zA-Z0-9\/+=,.@_-]*)\/\*$"
90-
S3_BUCKET: typing.ClassVar[str] = r"((s3:\/\/|s3a:\/\/)([a-zA-Z0-9+=,.@_-]*))\/.*$"
90+
S3_BUCKET: typing.ClassVar[str] = r"((s3:\/\/|s3a:\/\/)([a-zA-Z0-9+=,.@_-]*))(\/.*$)?"
9191
S3_PREFIX: typing.ClassVar[str] = "arn:aws:s3:::"
9292
S3_PATH_REGEX: typing.ClassVar[str] = r"((s3:\/\/)|(s3a:\/\/))(.*)"
9393
UC_MASTER_ROLES_ARN: typing.ClassVar[list[str]] = [

src/databricks/labs/ucx/aws/access.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,8 +213,9 @@ def _identify_missing_paths(self):
213213
missing_paths = set()
214214
for external_location in external_locations:
215215
matching_role = False
216+
path = PurePath(external_location.location)
216217
for role in compatible_roles:
217-
if external_location.location.startswith(role.resource_path):
218+
if PurePath(role.resource_path) in path.parents or path.match(role.resource_path):
218219
matching_role = True
219220
continue
220221
if matching_role:

tests/unit/aws/test_access.py

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
GetWorkspaceWarehouseConfigResponseSecurityPolicy,
2424
)
2525

26-
from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole
26+
from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole, AWSUCRoleCandidate
2727
from databricks.labs.ucx.aws.access import AWSResourcePermissions
2828
from databricks.labs.ucx.aws.credentials import IamRoleCreation
2929
from databricks.labs.ucx.aws.locations import AWSExternalLocationsMigration
@@ -103,7 +103,12 @@ def installation_no_roles():
103103
@pytest.fixture
104104
def backend():
105105
rows = {
106-
"external_locations": [["s3://BUCKET1/FOLDER1", 1], ["s3://BUCKET2/FOLDER2", 1], ["s3://BUCKETX/FOLDERX", 1]]
106+
"external_locations": [
107+
["s3://BUCKET1/FOLDER1", 1],
108+
["s3://BUCKET2/FOLDER2", 1],
109+
["s3://BUCKET4", 2],
110+
["s3://BUCKETX/FOLDERX", 1],
111+
]
107112
}
108113
return MockBackend(rows=rows, fails_on_first={})
109114

@@ -201,6 +206,25 @@ def test_create_external_locations_skip_existing(mock_ws, backend, locations):
201206
principal_acl.apply_location_acl.assert_called()
202207

203208

209+
def test_uc_roles_create_all_roles(mock_ws, backend, locations):
210+
install = MockInstallation({"uc_roles_access.csv": []})
211+
mock_ws.storage_credentials.list.return_value = []
212+
mock_ws.external_locations.list.return_value = []
213+
aws = AWSResources("profile", lambda cmd: (0, '{"Role": {"Arn": "arn:aws:iam::12345:role/role1"}}', ""))
214+
aws_resource_permissions = AWSResourcePermissions(install, mock_ws, aws, locations)
215+
216+
roles = aws_resource_permissions.list_uc_roles(single_role=False)
217+
expected_roles = [
218+
AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET4', policy_name='UC_POLICY', resource_paths=['s3://BUCKET4']),
219+
AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET1', policy_name='UC_POLICY', resource_paths=['s3://BUCKET1']),
220+
AWSUCRoleCandidate(role_name='UC_ROLE_BUCKETX', policy_name='UC_POLICY', resource_paths=['s3://BUCKETX']),
221+
AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET2', policy_name='UC_POLICY', resource_paths=['s3://BUCKET2']),
222+
]
223+
assert len(roles) == len(expected_roles)
224+
for role in expected_roles:
225+
assert role in roles
226+
227+
204228
def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installation, backend, locations):
205229
instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
206230
cluster_policy = Policy(
@@ -228,7 +252,7 @@ def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installatio
228252
aws.put_role_policy.assert_called_with(
229253
'UCX_MIGRATION_ROLE_ucx',
230254
'UCX_MIGRATION_POLICY_ucx',
231-
{'s3://BUCKET1/FOLDER1', 's3://BUCKET2/FOLDER2', 's3://BUCKETX/FOLDERX'},
255+
{'s3://BUCKET1/FOLDER1', 's3://BUCKET2/FOLDER2', 's3://BUCKET4', 's3://BUCKETX/FOLDERX'},
232256
None,
233257
None,
234258
)
@@ -416,7 +440,7 @@ def test_create_uc_role_single(mock_ws, installation_single_role, backend, locat
416440
call(
417441
'UC_ROLE_123123',
418442
'UC_POLICY',
419-
{'s3://BUCKET1', 's3://BUCKET1/*', 's3://BUCKET2', 's3://BUCKET2/*'},
443+
{'s3://BUCKET2/*', 's3://BUCKET4/*', 's3://BUCKET2', 's3://BUCKET4', 's3://BUCKET1', 's3://BUCKET1/*'},
420444
None,
421445
None,
422446
)
@@ -452,8 +476,6 @@ def test_create_uc_role_multiple_raises_error(mock_ws, installation_single_role,
452476
aws.list_all_uc_roles.return_value = []
453477
with pytest.raises(PermissionDenied):
454478
role_creation.run(MockPrompts({"Above *": "yes"}), single_role=False)
455-
assert call('UC_ROLE_BUCKET1') in aws.create_uc_role.call_args_list
456-
assert call('UC_ROLE_BUCKET2') in aws.create_uc_role.call_args_list
457479
aws.delete_role.assert_called_once()
458480

459481

0 commit comments

Comments
 (0)