From 1999774183b25d865d4c6ef1c7685f6491279f77 Mon Sep 17 00:00:00 2001 From: Liran Bareket Date: Fri, 10 Jan 2025 12:14:20 -0500 Subject: [PATCH 1/3] Addressed Bug and added tests --- src/databricks/labs/ucx/assessment/aws.py | 2 +- src/databricks/labs/ucx/aws/access.py | 3 ++- tests/unit/aws/test_access.py | 30 ++++++++++++++++++++--- 3 files changed, 30 insertions(+), 5 deletions(-) diff --git a/src/databricks/labs/ucx/assessment/aws.py b/src/databricks/labs/ucx/assessment/aws.py index 616574fecc..06df883c69 100644 --- a/src/databricks/labs/ucx/assessment/aws.py +++ b/src/databricks/labs/ucx/assessment/aws.py @@ -87,7 +87,7 @@ class AWSResources: S3_ACTIONS: typing.ClassVar[set[str]] = {"s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:PutObjectAcl"} S3_READONLY: typing.ClassVar[str] = "s3:GetObject" S3_REGEX: typing.ClassVar[str] = r"arn:aws:s3:::([a-zA-Z0-9\/+=,.@_-]*)\/\*$" - S3_BUCKET: typing.ClassVar[str] = r"((s3:\/\/|s3a:\/\/)([a-zA-Z0-9+=,.@_-]*))\/.*$" + S3_BUCKET: typing.ClassVar[str] = r"((s3:\/\/|s3a:\/\/)([a-zA-Z0-9+=,.@_-]*))(\/.*$)?" S3_PREFIX: typing.ClassVar[str] = "arn:aws:s3:::" S3_PATH_REGEX: typing.ClassVar[str] = r"((s3:\/\/)|(s3a:\/\/))(.*)" UC_MASTER_ROLES_ARN: typing.ClassVar[list[str]] = [ diff --git a/src/databricks/labs/ucx/aws/access.py b/src/databricks/labs/ucx/aws/access.py index 539471feed..946ba10ef6 100644 --- a/src/databricks/labs/ucx/aws/access.py +++ b/src/databricks/labs/ucx/aws/access.py @@ -213,8 +213,9 @@ def _identify_missing_paths(self): missing_paths = set() for external_location in external_locations: matching_role = False + path = PurePath(external_location.location) for role in compatible_roles: - if external_location.location.startswith(role.resource_path): + if PurePath(role.resource_path) in path.parents or path.match(role.resource_path): matching_role = True continue if matching_role: diff --git a/tests/unit/aws/test_access.py b/tests/unit/aws/test_access.py index 127185a893..803c15c29d 100644 --- a/tests/unit/aws/test_access.py +++ b/tests/unit/aws/test_access.py @@ -23,7 +23,7 @@ GetWorkspaceWarehouseConfigResponseSecurityPolicy, ) -from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole +from databricks.labs.ucx.assessment.aws import AWSPolicyAction, AWSResources, AWSRole, AWSUCRoleCandidate from databricks.labs.ucx.aws.access import AWSResourcePermissions from databricks.labs.ucx.aws.credentials import IamRoleCreation from databricks.labs.ucx.aws.locations import AWSExternalLocationsMigration @@ -103,7 +103,12 @@ def installation_no_roles(): @pytest.fixture def backend(): rows = { - "external_locations": [["s3://BUCKET1/FOLDER1", 1], ["s3://BUCKET2/FOLDER2", 1], ["s3://BUCKETX/FOLDERX", 1]] + "external_locations": [ + ["s3://BUCKET1/FOLDER1", 1], + ["s3://BUCKET2/FOLDER2", 1], + ["s3://BUCKET4", 2], + ["s3://BUCKETX/FOLDERX", 1], + ] } return MockBackend(rows=rows, fails_on_first={}) @@ -201,6 +206,25 @@ def test_create_external_locations_skip_existing(mock_ws, backend, locations): principal_acl.apply_location_acl.assert_called() +def test_uc_roles_create_all_roles(mock_ws, backend, locations): + install = MockInstallation({"uc_roles_access.csv": []}) + mock_ws.storage_credentials.list.return_value = [] + mock_ws.external_locations.list.return_value = [] + aws = AWSResources("profile", lambda cmd: (0, '{"Role": {"Arn": "arn:aws:iam::12345:role/role1"}}', "")) + aws_resource_permissions = AWSResourcePermissions(install, mock_ws, aws, locations) + + roles = aws_resource_permissions.list_uc_roles(single_role=False) + expected_roles = [ + AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET4', policy_name='UC_POLICY', resource_paths=['s3://BUCKET4']), + AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET1', policy_name='UC_POLICY', resource_paths=['s3://BUCKET1']), + AWSUCRoleCandidate(role_name='UC_ROLE_BUCKETX', policy_name='UC_POLICY', resource_paths=['s3://BUCKETX']), + AWSUCRoleCandidate(role_name='UC_ROLE_BUCKET2', policy_name='UC_POLICY', resource_paths=['s3://BUCKET2']), + ] + assert len(roles) == len(expected_roles) + for role in expected_roles: + assert role in roles + + def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installation, backend, locations): instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1" cluster_policy = Policy( @@ -228,7 +252,7 @@ def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installatio aws.put_role_policy.assert_called_with( 'UCX_MIGRATION_ROLE_ucx', 'UCX_MIGRATION_POLICY_ucx', - {'s3://BUCKET1/FOLDER1', 's3://BUCKET2/FOLDER2', 's3://BUCKETX/FOLDERX'}, + {'s3://BUCKET1/FOLDER1', 's3://BUCKET2/FOLDER2', 's3://BUCKET4', 's3://BUCKETX/FOLDERX'}, None, None, ) From 7c5862d220586d6325455020c2920d51a799ba8e Mon Sep 17 00:00:00 2001 From: Liran Bareket Date: Mon, 13 Jan 2025 09:46:22 -0500 Subject: [PATCH 2/3] Addressed Bug and added tests --- tests/unit/aws/test_access.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/aws/test_access.py b/tests/unit/aws/test_access.py index 803c15c29d..7fa250e345 100644 --- a/tests/unit/aws/test_access.py +++ b/tests/unit/aws/test_access.py @@ -440,7 +440,7 @@ def test_create_uc_role_single(mock_ws, installation_single_role, backend, locat call( 'UC_ROLE_123123', 'UC_POLICY', - {'s3://BUCKET1', 's3://BUCKET1/*', 's3://BUCKET2', 's3://BUCKET2/*'}, + {'s3://BUCKET2/*', 's3://BUCKET4/*', 's3://BUCKET2', 's3://BUCKET4', 's3://BUCKET1', 's3://BUCKET1/*'}, None, None, ) @@ -476,7 +476,7 @@ def test_create_uc_role_multiple_raises_error(mock_ws, installation_single_role, aws.list_all_uc_roles.return_value = [] with pytest.raises(PermissionDenied): role_creation.run(MockPrompts({"Above *": "yes"}), single_role=False) - assert call('UC_ROLE_BUCKET1') in aws.create_uc_role.call_args_list + assert call('UC_ROLE_BUCKET4') in aws.create_uc_role.call_args_list assert call('UC_ROLE_BUCKET2') in aws.create_uc_role.call_args_list aws.delete_role.assert_called_once() From 72b1ab665b44cb400e994c0744ed2a1eb5edd1ce Mon Sep 17 00:00:00 2001 From: Liran Bareket Date: Mon, 13 Jan 2025 10:10:43 -0500 Subject: [PATCH 3/3] Fixed issue with unit test --- tests/unit/aws/test_access.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unit/aws/test_access.py b/tests/unit/aws/test_access.py index 7fa250e345..4f856420c4 100644 --- a/tests/unit/aws/test_access.py +++ b/tests/unit/aws/test_access.py @@ -476,8 +476,6 @@ def test_create_uc_role_multiple_raises_error(mock_ws, installation_single_role, aws.list_all_uc_roles.return_value = [] with pytest.raises(PermissionDenied): role_creation.run(MockPrompts({"Above *": "yes"}), single_role=False) - assert call('UC_ROLE_BUCKET4') in aws.create_uc_role.call_args_list - assert call('UC_ROLE_BUCKET2') in aws.create_uc_role.call_args_list aws.delete_role.assert_called_once()