From 564e504b6bf9977d2398be84dae46ec3672ec39b Mon Sep 17 00:00:00 2001 From: Hari Selvarajan <105197202+HariGS-DB@users.noreply.github.com> Date: Wed, 25 Sep 2024 12:59:03 +0100 Subject: [PATCH] Fixes issue of circular dependency of migrate-location ACL (#2741) ## Changes Resolves #2529 ### Functionality - [ ] modified existing command: `databricks labs ucx migrate-locations` ### Tests - [ ] added unit tests - [ ] added integration tests --------- Co-authored-by: Amin Movahed Co-authored-by: Amin Movahed <140028681+aminmovahed-db@users.noreply.github.com> --- .../labs/ucx/contexts/application.py | 20 ++++++++----------- .../labs/ucx/hive_metastore/grants.py | 6 +++--- tests/unit/azure/test_locations.py | 2 +- .../hive_metastore/test_principal_grants.py | 14 ++++++++----- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/src/databricks/labs/ucx/contexts/application.py b/src/databricks/labs/ucx/contexts/application.py index 3bf6e8ddc3..95944a3d2a 100644 --- a/src/databricks/labs/ucx/contexts/application.py +++ b/src/databricks/labs/ucx/contexts/application.py @@ -18,7 +18,6 @@ from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler from databricks.labs.ucx.source_code.python_libraries import PythonLibraryResolver from databricks.sdk import AccountClient, WorkspaceClient, core -from databricks.sdk.errors import NotFound from databricks.sdk.service import sql from databricks.labs.ucx.account.workspaces import WorkspaceInfo @@ -305,18 +304,15 @@ def aws_acl(self): ) @cached_property - def principal_locations(self): - eligible_locations = {} - try: + def principal_locations_retriever(self): + def inner(): if self.is_azure: - eligible_locations = self.azure_acl.get_eligible_locations_principals() + return self.azure_acl.get_eligible_locations_principals() if self.is_aws: - eligible_locations = self.aws_acl.get_eligible_locations_principals() - if self.is_gcp: - raise NotImplementedError("Not implemented for GCP.") - except NotFound: - pass - return eligible_locations + return self.aws_acl.get_eligible_locations_principals() + raise NotImplementedError("Not implemented for GCP.") + + return inner @cached_property def principal_acl(self): @@ -326,7 +322,7 @@ def principal_acl(self): self.installation, self.tables_crawler, self.mounts_crawler, - self.principal_locations, + self.principal_locations_retriever, ) @cached_property diff --git a/src/databricks/labs/ucx/hive_metastore/grants.py b/src/databricks/labs/ucx/hive_metastore/grants.py index a2afa473f3..8673779697 100644 --- a/src/databricks/labs/ucx/hive_metastore/grants.py +++ b/src/databricks/labs/ucx/hive_metastore/grants.py @@ -579,7 +579,7 @@ def __init__( installation: Installation, tables_crawler: TablesCrawler, mounts_crawler: Mounts, - cluster_locations: list[ComputeLocations], + cluster_locations: Callable[[], list[ComputeLocations]], ): self._backend = backend self._ws = ws @@ -593,7 +593,7 @@ def get_interactive_cluster_grants(self) -> list[Grant]: mounts = list(self._mounts_crawler.snapshot()) grants: set[Grant] = set() - for compute_location in self._compute_locations: + for compute_location in self._compute_locations(): principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type) if len(principals) == 0: continue @@ -697,7 +697,7 @@ def apply_location_acl(self): "CREATE EXTERNAL VOLUME and READ_FILES for existing eligible interactive cluster users" ) # get the eligible location mapped for each interactive cluster - for compute_location in self._compute_locations: + for compute_location in self._compute_locations(): # get interactive cluster users principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type) if len(principals) == 0: diff --git a/tests/unit/azure/test_locations.py b/tests/unit/azure/test_locations.py index 2b6e8a177a..f1b901638b 100644 --- a/tests/unit/azure/test_locations.py +++ b/tests/unit/azure/test_locations.py @@ -30,7 +30,7 @@ def location_migration_for_test(ws, mock_backend, mock_installation, azurerm=Non azure_resource_permissions = AzureResourcePermissions(mock_installation, ws, azurerm, location_crawler) tables_crawler = TablesCrawler(mock_backend, 'ucx') mounts_crawler = Mounts(mock_backend, ws, 'ucx') - principal_acl = PrincipalACL(ws, mock_backend, mock_installation, tables_crawler, mounts_crawler, {}) + principal_acl = PrincipalACL(ws, mock_backend, mock_installation, tables_crawler, mounts_crawler, lambda: []) external_locations_migration = ExternalLocationsMigration( ws, location_crawler, azure_resource_permissions, azurerm, principal_acl ) diff --git a/tests/unit/hive_metastore/test_principal_grants.py b/tests/unit/hive_metastore/test_principal_grants.py index e62a998905..bef584fdab 100644 --- a/tests/unit/hive_metastore/test_principal_grants.py +++ b/tests/unit/hive_metastore/test_principal_grants.py @@ -157,11 +157,15 @@ def principal_acl(w, install, cluster_spn: list, warehouse_spn: list): spn_crawler = create_autospec(AzureServicePrincipalCrawler) spn_crawler.get_cluster_to_storage_mapping.return_value = cluster_spn - locations = [] - if w.config.is_azure: - locations = azure_acl(w, install, cluster_spn, warehouse_spn).get_eligible_locations_principals() - if w.config.is_aws: - locations = aws_acl(w, install).get_eligible_locations_principals() + + def inner(): + if w.config.is_azure: + return azure_acl(w, install, cluster_spn, warehouse_spn).get_eligible_locations_principals() + if w.config.is_aws: + return aws_acl(w, install).get_eligible_locations_principals() + return None + + locations = inner return PrincipalACL(w, sql_backend, install, table_crawler, mount_crawler, locations)