From 0820221b73ff15221920b3866f3b6c985ce2ff93 Mon Sep 17 00:00:00 2001 From: FastLee Date: Sun, 19 May 2024 15:12:34 -0400 Subject: [PATCH 1/2] Getting aws account id from the profile. --- src/databricks/labs/ucx/contexts/workspace_cli.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/databricks/labs/ucx/contexts/workspace_cli.py b/src/databricks/labs/ucx/contexts/workspace_cli.py index b6a7a207e9..97b00531f3 100644 --- a/src/databricks/labs/ucx/contexts/workspace_cli.py +++ b/src/databricks/labs/ucx/contexts/workspace_cli.py @@ -146,12 +146,15 @@ def aws_resources(self): @cached_property def aws_resource_permissions(self): + aws_account_id = self.named_parameters.get("aws_account_id") + if not aws_account_id: + aws_account_id = self.aws_resources.validate_connection().get("Account") return AWSResourcePermissions( self.installation, self.workspace_client, self.aws_resources, self.external_locations, - self.named_parameters.get("aws_account_id"), + aws_account_id, self.named_parameters.get("kms_key"), ) From a66b872eddbcb56353838058ea99654c3a7cb38c Mon Sep 17 00:00:00 2001 From: FastLee Date: Sun, 19 May 2024 15:51:18 -0400 Subject: [PATCH 2/2] Fixed Unit Test. --- tests/unit/test_cli.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index d4e9526d28..5319e4e036 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -388,7 +388,9 @@ def test_migrate_locations_azure(ws): def test_migrate_locations_aws(ws, caplog): - ctx = WorkspaceContext(ws).replace(is_aws=True, is_azure=False, aws_profile="profile") + ctx = WorkspaceContext(ws).replace( + is_aws=True, is_azure=False, aws_profile="profile", named_parameters={"aws_account_id": "123456789012"} + ) migrate_locations(ws, ctx=ctx) ws.external_locations.list.assert_called()