Skip to content

Commit

Permalink
Fixes issue of circular dependency of migrate-location ACL (databrick…
Browse files Browse the repository at this point in the history
…slabs#2741)

<!-- REMOVE IRRELEVANT COMMENTS BEFORE CREATING A PULL REQUEST -->
## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->


Resolves databrickslabs#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>
  • Loading branch information
3 people authored and jgarciaf106 committed Sep 26, 2024
1 parent 21da7cc commit 2a09a8f
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 21 deletions.
20 changes: 8 additions & 12 deletions src/databricks/labs/ucx/contexts/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -310,18 +309,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):
Expand All @@ -331,7 +327,7 @@ def principal_acl(self):
self.installation,
self.tables_crawler,
self.mounts_crawler,
self.principal_locations,
self.principal_locations_retriever,
)

@cached_property
Expand Down
6 changes: 3 additions & 3 deletions src/databricks/labs/ucx/hive_metastore/grants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/azure/test_locations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
14 changes: 9 additions & 5 deletions tests/unit/hive_metastore/test_principal_grants.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down

0 comments on commit 2a09a8f

Please sign in to comment.