diff --git a/src/databricks/labs/ucx/assessment/aws.py b/src/databricks/labs/ucx/assessment/aws.py index 2400fda4bd..2e10e03fa2 100644 --- a/src/databricks/labs/ucx/assessment/aws.py +++ b/src/databricks/labs/ucx/assessment/aws.py @@ -257,14 +257,13 @@ def _glue_policy_actions(self, action): actions = action.get("Action") if not actions: return [] - if isinstance(actions, list): - if "glue:*" not in actions: - # Check if all the required glue action are present in the role - for required_action in self.GLUE_REQUIRED_ACTIONS: - if required_action not in actions: - return [] - elif actions != "glue:*": - return [] + if not isinstance(actions, list): + actions = [actions] + if "glue:*" not in actions: + # Check if all the required glue action are present in the role + for required_action in self.GLUE_REQUIRED_ACTIONS: + if required_action not in actions: + return [] if "*" not in action.get("Resource", []): return [] diff --git a/src/databricks/labs/ucx/aws/credentials.py b/src/databricks/labs/ucx/aws/credentials.py index 67192929c8..a491114053 100644 --- a/src/databricks/labs/ucx/aws/credentials.py +++ b/src/databricks/labs/ucx/aws/credentials.py @@ -91,8 +91,12 @@ def create_service_credential(self, name: str, role_arn: str) -> str | None: "skip_validation": True, }, ) - if response and isinstance(response, dict): - return response.get("aws_iam_role").get("external_id") + if not response or not isinstance(response, dict): + return None + iam_role = response.get("aws_iam_role") + if not iam_role or not isinstance(iam_role, dict): + return None + return iam_role.get("external_id") def validate(self, role_action: AWSRoleAction) -> CredentialValidationResult: try: diff --git a/tests/unit/aws/test_access.py b/tests/unit/aws/test_access.py index 9e0e1112b1..89f2af40aa 100644 --- a/tests/unit/aws/test_access.py +++ b/tests/unit/aws/test_access.py @@ -67,6 +67,32 @@ def installation_single_role(): ) +@pytest.fixture +def installation_glue(): + return MockInstallation( + { + "config.yml": { + 'version': 2, + 'inventory_database': 'ucx', + 'policy_id': 'policy_id', + 'spark_conf': {'spark.databricks.hive.metastore.glueCatalog.enabled': 'true'}, + 'connect': { + 'host': 'foo', + 'token': 'bar', + }, + }, + "uc_roles_access.csv": [ + { + "role_arn": "arn:aws:iam::12345:role/uc-role1", + "resource_type": "s3", + "privilege": "WRITE_FILES", + "resource_path": "s3://BUCKETX", + } + ], + } + ) + + @pytest.fixture def installation_multiple_roles(): return MockInstallation( @@ -467,6 +493,29 @@ def test_create_uc_role_single(mock_ws, installation_single_role, backend, locat ) +def test_create_uc_role_glue(mock_ws, installation_glue, backend, locations): + mock_ws.metastores.current.return_value = MetastoreAssignment(metastore_id="123123", workspace_id="456456") + aws = create_autospec(AWSResources) + aws.validate_connection.return_value = {} + config = installation_glue.load(WorkspaceConfig) + aws_resource_permissions = AWSResourcePermissions(installation_glue, mock_ws, config, aws, locations) + role_creation = IamRoleCreation(installation_glue, mock_ws, aws_resource_permissions) + aws.list_all_uc_roles.return_value = [] + role_creation.run(MockPrompts({"Above *": "yes"}), single_role=True) + assert aws.create_uc_role.assert_called + assert ( + call( + 'UC_ROLE_123123', + 'UC_POLICY', + 'glue', + {'*'}, + None, + None, + ) + in aws.put_role_policy.call_args_list + ) + + def test_create_uc_role_multiple(mock_ws, installation_single_role, backend, locations): aws = create_autospec(AWSResources) aws.validate_connection.return_value = {} @@ -622,6 +671,86 @@ def test_get_uc_compatible_roles(mock_ws, mock_installation, locations): ) +def test_get_uc_compatible_roles_glue(mock_ws, mock_installation, locations): + aws = create_autospec(AWSResources) + aws.get_role_policy.side_effect = [ + [ + AWSPolicyAction( + resource_type="s3", + privilege="READ_FILES", + resource_path="s3://bucket1", + ), + AWSPolicyAction( + resource_type="s3", + privilege="READ_FILES", + resource_path="s3://bucket2", + ), + AWSPolicyAction( + resource_type="s3", + privilege="READ_FILES", + resource_path="s3://bucket3", + ), + ], + [], + [], + [ + AWSPolicyAction( + resource_type="glue", + privilege="USAGE", + resource_path="*", + ), + ], + [], + [], + ] + aws.list_role_policies.return_value = ["Policy1", "Policy2", "Policy3"] + aws.list_attached_policies_in_role.return_value = [ + "arn:aws:iam::aws:policy/Policy1", + "arn:aws:iam::aws:policy/Policy2", + ] + aws.list_all_uc_roles.return_value = [ + AWSRole(path='/', role_name='uc-role1', role_id='12345', arn='arn:aws:iam::12345:role/uc-role1') + ] + config = mock_installation.load(WorkspaceConfig) + aws_resource_permissions = AWSResourcePermissions( + mock_installation, + mock_ws, + config, + aws, + locations, + ) + aws_resource_permissions.save_uc_compatible_roles() + mock_installation.assert_file_written( + 'uc_roles_access.csv', + [ + { + 'privilege': 'READ_FILES', + 'resource_path': 's3://bucket1', + 'resource_type': 's3', + 'role_arn': 'arn:aws:iam::12345:role/uc-role1', + }, + { + 'privilege': 'READ_FILES', + 'resource_path': 's3://bucket2', + 'resource_type': 's3', + 'role_arn': 'arn:aws:iam::12345:role/uc-role1', + }, + { + 'privilege': 'READ_FILES', + 'resource_path': 's3://bucket3', + 'resource_type': 's3', + 'role_arn': 'arn:aws:iam::12345:role/uc-role1', + }, + { + 'privilege': 'USAGE', + 'resource_path': '*', + 'resource_type': 'glue', + 'role_arn': 'arn:aws:iam::12345:role/uc-role1', + }, + ], + ) + + def test_instance_profiles_empty_mapping(mock_ws, mock_installation, locations, caplog): aws = create_autospec(AWSResources) aws.get_instance_profile_role_arn.return_value = "arn:aws:iam::12345:role/role1" diff --git a/tests/unit/aws/test_credentials.py b/tests/unit/aws/test_credentials.py index c6c44e625c..8c996b6730 100644 --- a/tests/unit/aws/test_credentials.py +++ b/tests/unit/aws/test_credentials.py @@ -156,6 +156,28 @@ def test_run_no_credential_to_migrate(caplog, installation, credential_manager): assert "No IAM Role to migrate" in caplog.messages +def test_create_glue_credentials(installation): + prompts = MockPrompts({"Above IAM roles will be migrated*": "Yes"}) + resource_permissions = create_autospec(AWSResourcePermissions) + resource_permissions.get_glue_roles.return_value = [ + AWSCredentialCandidate( + role_arn="arn:aws:iam::123456789012:role/glue-role", + privilege=Privilege.USAGE.value, + paths={"*"}, + ) + ] + storage_credentials = create_autospec(CredentialManager) + storage_credentials.create_service_credential.return_value = "1234" + migration = IamRoleMigration(installation, resource_permissions, storage_credentials) + migration.migrate_glue(prompts) + storage_credentials.create_service_credential.assert_called_with( + name="glue-role", role_arn="arn:aws:iam::123456789012:role/glue-role" + ) + resource_permissions.update_uc_role.assert_called_with( + "glue-role", "arn:aws:iam::123456789012:role/glue-role", "1234" + ) + + def test_validate_read_only_storage_credentials(credential_manager): role_action = AWSRoleAction("arn:aws:iam::123456789012:role/client_id", "s3", "READ_FILES", "s3://prefix")