Skip to content

Commit 564e504

Browse files
Fixes issue of circular dependency of migrate-location ACL (#2741)
<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST --> ## Changes <!-- Summary of your changes that are easy to understand. Add screenshots when necessary --> Resolves #2529 ### Functionality - [ ] modified existing command: `databricks labs ucx migrate-locations` ### Tests <!-- How is this tested? Please see the checklist below and also describe any other relevant tests --> - [ ] added unit tests - [ ] added integration tests --------- Co-authored-by: Amin Movahed <amin.movahed@databricks.com> Co-authored-by: Amin Movahed <140028681+aminmovahed-db@users.noreply.github.com>
1 parent 5b3c0f2 commit 564e504

File tree

4 files changed

+21
-21
lines changed

4 files changed

+21
-21
lines changed

src/databricks/labs/ucx/contexts/application.py

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler
1919
from databricks.labs.ucx.source_code.python_libraries import PythonLibraryResolver
2020
from databricks.sdk import AccountClient, WorkspaceClient, core
21-
from databricks.sdk.errors import NotFound
2221
from databricks.sdk.service import sql
2322

2423
from databricks.labs.ucx.account.workspaces import WorkspaceInfo
@@ -305,18 +304,15 @@ def aws_acl(self):
305304
)
306305

307306
@cached_property
308-
def principal_locations(self):
309-
eligible_locations = {}
310-
try:
307+
def principal_locations_retriever(self):
308+
def inner():
311309
if self.is_azure:
312-
eligible_locations = self.azure_acl.get_eligible_locations_principals()
310+
return self.azure_acl.get_eligible_locations_principals()
313311
if self.is_aws:
314-
eligible_locations = self.aws_acl.get_eligible_locations_principals()
315-
if self.is_gcp:
316-
raise NotImplementedError("Not implemented for GCP.")
317-
except NotFound:
318-
pass
319-
return eligible_locations
312+
return self.aws_acl.get_eligible_locations_principals()
313+
raise NotImplementedError("Not implemented for GCP.")
314+
315+
return inner
320316

321317
@cached_property
322318
def principal_acl(self):
@@ -326,7 +322,7 @@ def principal_acl(self):
326322
self.installation,
327323
self.tables_crawler,
328324
self.mounts_crawler,
329-
self.principal_locations,
325+
self.principal_locations_retriever,
330326
)
331327

332328
@cached_property

src/databricks/labs/ucx/hive_metastore/grants.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -579,7 +579,7 @@ def __init__(
579579
installation: Installation,
580580
tables_crawler: TablesCrawler,
581581
mounts_crawler: Mounts,
582-
cluster_locations: list[ComputeLocations],
582+
cluster_locations: Callable[[], list[ComputeLocations]],
583583
):
584584
self._backend = backend
585585
self._ws = ws
@@ -593,7 +593,7 @@ def get_interactive_cluster_grants(self) -> list[Grant]:
593593
mounts = list(self._mounts_crawler.snapshot())
594594
grants: set[Grant] = set()
595595

596-
for compute_location in self._compute_locations:
596+
for compute_location in self._compute_locations():
597597
principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type)
598598
if len(principals) == 0:
599599
continue
@@ -697,7 +697,7 @@ def apply_location_acl(self):
697697
"CREATE EXTERNAL VOLUME and READ_FILES for existing eligible interactive cluster users"
698698
)
699699
# get the eligible location mapped for each interactive cluster
700-
for compute_location in self._compute_locations:
700+
for compute_location in self._compute_locations():
701701
# get interactive cluster users
702702
principals = self._get_cluster_principal_mapping(compute_location.compute_id, compute_location.compute_type)
703703
if len(principals) == 0:

tests/unit/azure/test_locations.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def location_migration_for_test(ws, mock_backend, mock_installation, azurerm=Non
3030
azure_resource_permissions = AzureResourcePermissions(mock_installation, ws, azurerm, location_crawler)
3131
tables_crawler = TablesCrawler(mock_backend, 'ucx')
3232
mounts_crawler = Mounts(mock_backend, ws, 'ucx')
33-
principal_acl = PrincipalACL(ws, mock_backend, mock_installation, tables_crawler, mounts_crawler, {})
33+
principal_acl = PrincipalACL(ws, mock_backend, mock_installation, tables_crawler, mounts_crawler, lambda: [])
3434
external_locations_migration = ExternalLocationsMigration(
3535
ws, location_crawler, azure_resource_permissions, azurerm, principal_acl
3636
)

tests/unit/hive_metastore/test_principal_grants.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -157,11 +157,15 @@ def principal_acl(w, install, cluster_spn: list, warehouse_spn: list):
157157

158158
spn_crawler = create_autospec(AzureServicePrincipalCrawler)
159159
spn_crawler.get_cluster_to_storage_mapping.return_value = cluster_spn
160-
locations = []
161-
if w.config.is_azure:
162-
locations = azure_acl(w, install, cluster_spn, warehouse_spn).get_eligible_locations_principals()
163-
if w.config.is_aws:
164-
locations = aws_acl(w, install).get_eligible_locations_principals()
160+
161+
def inner():
162+
if w.config.is_azure:
163+
return azure_acl(w, install, cluster_spn, warehouse_spn).get_eligible_locations_principals()
164+
if w.config.is_aws:
165+
return aws_acl(w, install).get_eligible_locations_principals()
166+
return None
167+
168+
locations = inner
165169
return PrincipalACL(w, sql_backend, install, table_crawler, mount_crawler, locations)
166170

167171

0 commit comments

Comments
 (0)