Added AWS S3 support for migrate-locations command (#1009)

Author: FastLee

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

@@ -17,6 +17,7 @@
 
 from databricks.labs.ucx.framework.crawlers import StatementExecutionBackend
 from databricks.labs.ucx.hive_metastore import ExternalLocations
+from databricks.labs.ucx.hive_metastore.locations import ExternalLocation
 
 logger = logging.getLogger(__name__)
 
@@ -358,6 +359,11 @@ def get_uc_compatible_roles(self):
         return role_actions
 
     def create_uc_roles_cli(self, *, single_role=True, role_name="UC_ROLE", policy_name="UC_POLICY"):
+        # Get the missing paths
+        # Identify the S3 prefixes
+        # Create the roles and policies for the missing S3 prefixes
+        # If single_role is True, create a single role and policy for all the missing S3 prefixes
+        # If single_role is False, create a role and policy for each missing S3 prefix
         missing_paths = self._identify_missing_paths()
         s3_prefixes = set()
         for missing_path in missing_paths:
@@ -460,3 +466,64 @@ def save_instance_profile_permissions(self) -> str | None:
             logger.warning("No Mapping Was Generated.")
             return None
         return self._installation.save(instance_profile_access, filename=self.INSTANCE_PROFILES_FILE_NAMES)
+
+    def _identify_missing_external_locations(
+        self,
+        external_locations: Iterable[ExternalLocation],
+        existing_paths: list[str],
+        compatible_roles: list[AWSRoleAction],
+    ) -> set[tuple[str, str]]:
+        # Get recommended external locations
+        # Get existing external locations
+        # Get list of paths from get_uc_compatible_roles
+        # Identify recommended external location paths that don't have an external location and return them
+        missing_paths = set()
+        for external_location in external_locations:
+            existing = False
+            for path in existing_paths:
+                if path in external_location.location:
+                    existing = True
+                    continue
+            if existing:
+                continue
+            new_path = PurePath(external_location.location)
+            matching_role = None
+            for role in compatible_roles:
+                if new_path.match(role.resource_path):
+                    matching_role = role.role_arn
+                    continue
+            if matching_role:
+                missing_paths.add((external_location.location, matching_role))
+
+        return missing_paths
+
+    def _get_existing_credentials_dict(self):
+        credentials = self._ws.storage_credentials.list()
+        credentials_dict = {}
+        for credential in credentials:
+            credentials_dict[credential.aws_iam_role.role_arn] = credential.name
+        return credentials_dict
+
+    def create_external_locations(self, location_init="UCX_location"):
+        # For each path find out the role that has access to it
+        # Find out the credential that is pointing to this path
+        # Create external location for the path using the credential identified
+        credential_dict = self._get_existing_credentials_dict()
+        external_locations = ExternalLocations(self._ws, self._backend, self._schema).snapshot()
+        existing_external_locations = self._ws.external_locations.list()
+        existing_paths = [external_location.url for external_location in existing_external_locations]
+        compatible_roles = self.get_uc_compatible_roles()
+        missing_paths = self._identify_missing_external_locations(external_locations, existing_paths, compatible_roles)
+        external_location_names = [external_location.name for external_location in existing_external_locations]
+        external_location_num = 1
+        for path, role_arn in missing_paths:
+            if role_arn not in credential_dict:
+                logger.error(f"Missing credential for role {role_arn} for path {path}")
+                continue
+            while True:
+                external_location_name = f"{location_init}_{external_location_num}"
+                if external_location_name not in external_location_names:
+                    break
+                external_location_num += 1
+            self._ws.external_locations.create(external_location_name, path, credential_dict[role_arn])
+            external_location_num += 1

src/databricks/labs/ucx/cli.py

@@ -346,7 +346,7 @@ def create_uber_principal(w: WorkspaceClient, subscription_id: str):
 
 
 @ucx.command
-def migrate_locations(w: WorkspaceClient):
+def migrate_locations(w: WorkspaceClient, aws_profile: str | None = None):
     """This command creates UC external locations. The candidate locations to be created are extracted from guess_external_locations
     task in the assessment job. You can run validate_external_locations command to check the candidate locations. Please make sure
     the credentials haven migrated before running this command. The command will only create the locations that have corresponded UC Storage Credentials.
@@ -357,7 +357,15 @@ def migrate_locations(w: WorkspaceClient):
         service_principal_migration = ExternalLocationsMigration.for_cli(w, installation)
         service_principal_migration.run()
     if w.config.is_aws:
-        logger.error("migrate_locations is not yet supported in AWS")
+        logger.error("Migrate_locations for AWS")
+        if not shutil.which("aws"):
+            logger.error("Couldn't find AWS CLI in path. Please install the CLI from https://aws.amazon.com/cli/")
+            return
+        installation = Installation.current(w, 'ucx')
+        config = installation.load(WorkspaceConfig)
+        sql_backend = StatementExecutionBackend(w, config.warehouse_id)
+        aws_permissions = AWSResourcePermissions.for_cli(w, sql_backend, aws_profile, config.inventory_database)
+        aws_permissions.create_external_locations()
     if w.config.is_gcp:
         logger.error("migrate_locations is not yet supported in GCP")
 

tests/unit/assessment/test_aws.py

@@ -1,11 +1,17 @@
 import logging
+from unittest import mock
 from unittest.mock import MagicMock, call, create_autospec
 
 import pytest
 from databricks.labs.blueprint.installation import MockInstallation
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import ResourceDoesNotExist
 from databricks.sdk.service import iam
+from databricks.sdk.service.catalog import (
+    AwsIamRole,
+    ExternalLocationInfo,
+    StorageCredentialInfo,
+)
 from databricks.sdk.service.compute import InstanceProfile
 
 from databricks.labs.ucx.assessment.aws import (
@@ -973,3 +979,76 @@ def test_get_uc_compatible_roles():
             },
         ],
     )
+
+
+def test_create_external_locations():
+    ws = create_autospec(WorkspaceClient)
+    aws = create_autospec(AWSResources)
+    installation = MockInstallation()
+    installation.load = MagicMock()
+    installation.load.return_value = [
+        AWSRoleAction("arn:aws:iam::12345:role/uc-role1", "s3", "WRITE_FILES", "s3://BUCKET1/*"),
+        AWSRoleAction("arn:aws:iam::12345:role/uc-role1", "s3", "WRITE_FILES", "s3://BUCKET2/*"),
+        AWSRoleAction("arn:aws:iam::12345:role/uc-rolex", "s3", "WRITE_FILES", "s3://BUCKETX/FOLDERX"),
+    ]
+    rows = {
+        "external_locations": [["s3://BUCKET1/FOLDER1", 1], ["s3://BUCKET2/FOLDER2", 1], ["s3://BUCKETX/FOLDERX", 1]]
+    }
+    ws.storage_credentials.list.return_value = [
+        StorageCredentialInfo(
+            id="1",
+            name="cred1",
+            aws_iam_role=AwsIamRole("arn:aws:iam::12345:role/uc-role1"),
+        ),
+        StorageCredentialInfo(
+            id="2",
+            name="credx",
+            aws_iam_role=AwsIamRole("arn:aws:iam::12345:role/uc-rolex"),
+        ),
+    ]
+    errors = {}
+    backend = MockBackend(rows=rows, fails_on_first=errors)
+    aws_resource_permissions = AWSResourcePermissions(installation, ws, backend, aws, "ucx")
+    aws_resource_permissions.create_external_locations()
+    calls = [
+        call(mock.ANY, 's3://BUCKET1/FOLDER1', 'cred1'),
+        call(mock.ANY, 's3://BUCKET2/FOLDER2', 'cred1'),
+        call(mock.ANY, 's3://BUCKETX/FOLDERX', 'credx'),
+    ]
+    ws.external_locations.create.assert_has_calls(calls, any_order=True)
+
+
+def test_create_external_locations_skip_existing():
+    ws = create_autospec(WorkspaceClient)
+    aws = create_autospec(AWSResources)
+    installation = MockInstallation()
+    installation.load = MagicMock()
+    installation.load.return_value = [
+        AWSRoleAction("arn:aws:iam::12345:role/uc-role1", "s3", "WRITE_FILES", "s3://BUCKET1/*"),
+        AWSRoleAction("arn:aws:iam::12345:role/uc-rolex", "s3", "WRITE_FILES", "s3://BUCKETX/FOLDERX"),
+    ]
+    rows = {"external_locations": [["s3://BUCKET1/FOLDER1", 1], ["s3://BUCKETX/FOLDERX", 1]]}
+    ws.storage_credentials.list.return_value = [
+        StorageCredentialInfo(
+            id="1",
+            name="cred1",
+            aws_iam_role=AwsIamRole("arn:aws:iam::12345:role/uc-role1"),
+        ),
+        StorageCredentialInfo(
+            id="2",
+            name="credx",
+            aws_iam_role=AwsIamRole("arn:aws:iam::12345:role/uc-rolex"),
+        ),
+    ]
+    ws.external_locations.list.return_value = [
+        ExternalLocationInfo(name="UCX_FOO_1", url="s3://BUCKETX/FOLDERX", credential_name="credx"),
+    ]
+
+    errors = {}
+    backend = MockBackend(rows=rows, fails_on_first=errors)
+    aws_resource_permissions = AWSResourcePermissions(installation, ws, backend, aws, "ucx")
+    aws_resource_permissions.create_external_locations(location_init="UCX_FOO")
+    calls = [
+        call("UCX_FOO_2", 's3://BUCKET1/FOLDER1', 'cred1'),
+    ]
+    ws.external_locations.create.assert_has_calls(calls, any_order=True)

tests/unit/test_cli.py

@@ -368,12 +368,22 @@ def test_migrate_locations_azure(ws):
     ws.external_locations.list.assert_called()
 
 
-def test_migrate_locations_aws(ws, caplog):
+def test_migrate_locations_aws(ws, caplog, mocker):
+    mocker.patch("shutil.which", return_value="/path/aws")
     ws.config.is_azure = False
     ws.config.is_aws = True
     ws.config.is_gcp = False
-    migrate_locations(ws)
-    assert "migrate_locations is not yet supported in AWS" in caplog.messages
+    with pytest.raises(ResourceWarning):
+        migrate_locations(ws, aws_profile="profile")
+
+
+def test_missing_aws_cli(ws, caplog, mocker):
+    mocker.patch("shutil.which", return_value=None)
+    ws.config.is_azure = False
+    ws.config.is_aws = True
+    ws.config.is_gcp = False
+    migrate_locations(ws, aws_profile="profile")
+    assert "Couldn't find AWS CLI in path. Please install the CLI from https://aws.amazon.com/cli/" in caplog.messages
 
 
 def test_migrate_locations_gcp(ws, caplog):