Skip to content

Commit c1cb7a5

Browse files
authored
Group migration workflow improvements (#3391)
## Changes Group migration workflow improvements: - Return early when `include_databases` is provided - Fix unavailable external catalog on tacl cluster - Check outcome of verify permissions - Setup tacl cluster - Log what to do when group permissions did not verify successfully - Merge duplicate tests ### Functionality - [x] modified existing command: `databricks labs ucx group-migration-legacy` ### Tests - [x] modified integration tests
1 parent f0ead4b commit c1cb7a5

File tree

4 files changed

+96
-66
lines changed

4 files changed

+96
-66
lines changed

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -539,10 +539,10 @@ def _option_as_python(scala_option: typing.Any):
539539
return scala_option.get() if scala_option.isDefined() else None
540540

541541
def _all_databases(self) -> list[str]:
542-
try:
543-
if not self._include_database:
544-
return list(self._iterator(self._external_catalog.listDatabases()))
542+
if self._include_database:
545543
return self._include_database
544+
try:
545+
return list(self._iterator(self._external_catalog.listDatabases()))
546546
except Exception as err: # pylint: disable=broad-exception-caught
547547
if "py4j.security.Py4JSecurityException" in str(err):
548548
logger.error(

src/databricks/labs/ucx/workspace_access/manager.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,9 +116,15 @@ def verify_group_permissions(self) -> bool:
116116
verifier_tasks.extend(tasks_for_support)
117117

118118
logger.info(f"Starting to verify permissions. Total tasks: {len(verifier_tasks)}")
119-
Threads.strict("verify group permissions", verifier_tasks)
119+
verifications, errors = Threads.gather("verify group permissions", verifier_tasks)
120+
if errors:
121+
logger.error(f"Detected {len(errors)} errors while verifying permissions")
122+
raise ManyError(errors)
123+
if not all(verifications):
124+
unsuccessful_verifications = len(verifications) - sum(verifications)
125+
logger.warning(f"Detected {unsuccessful_verifications} out of {len(verifications)} invalidated permissions")
126+
return False
120127
logger.info("All permissions validated successfully. No issues found.")
121-
122128
return True
123129

124130
def object_type_support(self) -> dict[str, AclSupport]:

src/databricks/labs/ucx/workspace_access/workflows.py

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,12 @@ def reflect_account_groups_on_workspace(self, ctx: RuntimeContext):
3939
return
4040
ctx.group_manager.reflect_account_groups_on_workspace()
4141

42-
@job_task(depends_on=[reflect_account_groups_on_workspace], job_cluster="tacl")
42+
@job_task(job_cluster="tacl")
43+
def setup_tacl(self, ctx: RuntimeContext):
44+
"""(Optimization) Allow the TACL job cluster to be started while we're verifying the prerequisites for
45+
refreshing everything."""
46+
47+
@job_task(depends_on=[reflect_account_groups_on_workspace, setup_tacl], job_cluster="tacl")
4348
def apply_permissions_to_account_groups(self, ctx: RuntimeContext):
4449
"""Fourth phase of the workspace-local group migration process. It does the following:
4550
- Assigns the full set of permissions of the original group to the account-level one
@@ -65,7 +70,12 @@ def validate_groups_permissions(self, ctx: RuntimeContext):
6570
if not ctx.config.use_legacy_permission_migration:
6671
logger.info("Use `migrate-groups` job, or set `use_legacy_permission_migration: true` in config.yml.")
6772
return
68-
ctx.permission_manager.verify_group_permissions()
73+
if not ctx.permission_manager.verify_group_permissions():
74+
raise ValueError(
75+
"Some group permissions were not migrated successfully. Wait for an hour then use the "
76+
"`validate-group-permissions` workflow to validate the permissions after the API caught up. "
77+
f"Run `databricks labs ucx logs --workflow '{self._name}' --debug` for more details."
78+
)
6979

7080

7181
class PermissionsMigrationAPI(Workflow):
@@ -142,7 +152,10 @@ def validate_groups_permissions(self, ctx: RuntimeContext):
142152
if not ctx.config.use_legacy_permission_migration:
143153
logger.info("Use `migrate-groups` job, or set `use_legacy_permission_migration: true` in config.yml.")
144154
return
145-
ctx.permission_manager.verify_group_permissions()
155+
if not ctx.permission_manager.verify_group_permissions():
156+
raise ValueError(
157+
f"Some group permissions were not migrated successfully. Run `databricks labs ucx logs --workflow '{self._name}' --debug` for more details."
158+
)
146159

147160

148161
class RemoveWorkspaceLocalGroups(Workflow):

tests/integration/workspace_access/test_workflows.py

Lines changed: 69 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,105 @@
1+
import datetime as dt
12
from dataclasses import replace
23

3-
4+
import pytest
5+
from databricks.sdk.retries import retried
46
from databricks.sdk.service import sql
57
from databricks.sdk.service.iam import PermissionLevel
68
from databricks.sdk.service.workspace import AclPermission
79

8-
from databricks.labs.ucx.workspace_access.groups import MigratedGroup
9-
1010

1111
def test_running_real_migrate_groups_job(
1212
installation_ctx,
1313
make_cluster_policy,
1414
make_cluster_policy_permissions,
1515
make_secret_scope,
1616
make_secret_scope_acl,
17-
):
18-
ws_group_a, acc_group_a = installation_ctx.make_ucx_group(wait_for_provisioning=True)
19-
17+
) -> None:
18+
"""Test the migrate groups workflow.
19+
20+
We have many asserts in a single integration tests as we minimize the number of integration tests that run workflows
21+
to minimize the number of long-running integration tests.
22+
23+
For testing, we require:
24+
- UCX installation (as always)
25+
- A workspace group as the source for migration
26+
- A account group as the target for migration
27+
- Permissions to migrate:
28+
- Cluster policy
29+
- Schema and table access permissions
30+
- Secret scope permissions
31+
- Inventory tables used by the migrate groups workflow
32+
33+
We test:
34+
- The workflow to complete successfully
35+
- The workspace group to be renamed
36+
- The permissions to be transferred to the account group
37+
"""
38+
ws_group, acc_group = installation_ctx.make_ucx_group(wait_for_provisioning=True)
39+
40+
# TODO: Move `make_cluster_policy` and `make_cluster_policy_permissions` to context like other `make_` methods
2041
cluster_policy = make_cluster_policy()
2142
make_cluster_policy_permissions(
2243
object_id=cluster_policy.policy_id,
2344
permission_level=PermissionLevel.CAN_USE,
24-
group_name=ws_group_a.display_name,
45+
group_name=ws_group.display_name,
2546
)
2647

27-
table = installation_ctx.make_table()
28-
installation_ctx.make_grant(ws_group_a.display_name, "SELECT", table_info=table)
48+
schema = installation_ctx.make_schema()
49+
table = installation_ctx.make_table(schema_name=schema.name)
50+
installation_ctx.make_grant(ws_group.display_name, 'USAGE', schema_info=schema)
51+
installation_ctx.make_grant(ws_group.display_name, 'OWN', schema_info=schema)
52+
installation_ctx.make_grant(ws_group.display_name, 'SELECT', table_info=table)
2953

54+
# TODO: Move `make_secret_scope` and `make_secret_scope_acl` to context like other `make_` methods
3055
secret_scope = make_secret_scope()
31-
make_secret_scope_acl(scope=secret_scope, principal=ws_group_a.display_name, permission=AclPermission.WRITE)
56+
make_secret_scope_acl(scope=secret_scope, principal=ws_group.display_name, permission=AclPermission.WRITE)
3257

33-
installation_ctx.__dict__['include_group_names'] = [ws_group_a.display_name]
58+
# TODO: Move `include_object_permissions` to context like other `include_` attributes
59+
# Limit the considered permissions to the following objects:
3460
installation_ctx.__dict__['include_object_permissions'] = [
3561
f"cluster-policies:{cluster_policy.policy_id}",
3662
f"TABLE:{table.full_name}",
3763
f"secrets:{secret_scope}",
3864
]
3965

4066
installation_ctx.workspace_installation.run()
41-
42-
installation_ctx.deployed_workflows.run_workflow("migrate-groups")
67+
# The crawlers should run as part of the assessment. To minimize the crawling here, we only crawl what is necessary
68+
# Tables crawler fails on `tacl` cluster used by the apply and validate permission tasks
69+
installation_ctx.tables_crawler.snapshot(force_refresh=True)
70+
71+
workflow = "migrate-groups"
72+
installation_ctx.deployed_workflows.run_workflow(workflow, skip_job_wait=True)
73+
assert installation_ctx.deployed_workflows.validate_step(workflow), f"Workflow failed: {workflow}"
74+
75+
# Wrapper functions to wait for eventual consistency of API
76+
@retried(on=[KeyError], timeout=dt.timedelta(minutes=1))
77+
def wait_for_workspace_group_to_exists(display_name: str) -> bool:
78+
if installation_ctx.group_manager.has_workspace_group(display_name):
79+
return True
80+
raise KeyError(f"Group not found {display_name}")
81+
82+
# The original workspace group should be renamed
83+
renamed_workspace_group_name = installation_ctx.renamed_group_prefix + ws_group.display_name
84+
assert wait_for_workspace_group_to_exists(
85+
renamed_workspace_group_name
86+
), f"Workspace group not found: {renamed_workspace_group_name}"
87+
if installation_ctx.group_manager.has_workspace_group(ws_group.display_name): # Avoid wait on timeout
88+
with pytest.raises(TimeoutError):
89+
wait_for_workspace_group_to_exists(ws_group.display_name) # Expect to NOT exists
90+
91+
schema_grants = installation_ctx.grants_crawler.for_schema_info(schema)
92+
assert {"USAGE", "OWN"} == schema_grants[acc_group.display_name], "Incorrect schema grants for migrated group"
4393

4494
# specific permissions api migrations are checked in different and smaller integration tests
45-
found = installation_ctx.generic_permissions_support.load_as_dict("cluster-policies", cluster_policy.policy_id)
46-
assert acc_group_a.display_name in found, "Group not found in cluster policies"
47-
assert found[acc_group_a.display_name] == PermissionLevel.CAN_USE
95+
object_permissions = installation_ctx.generic_permissions_support.load_as_dict(
96+
"cluster-policies", cluster_policy.policy_id
97+
)
98+
assert acc_group.display_name in object_permissions, "Group not found in cluster policies"
99+
assert object_permissions[acc_group.display_name] == PermissionLevel.CAN_USE
48100

49101
scope_permission = installation_ctx.secret_scope_acl_support.secret_scope_permission(
50-
secret_scope, acc_group_a.display_name
102+
secret_scope, acc_group.display_name
51103
)
52104
assert scope_permission == AclPermission.WRITE
53105

@@ -96,44 +148,3 @@ def test_running_legacy_validate_groups_permissions_job(
96148

97149
# assert the job does not throw any exception
98150
installation_ctx.deployed_workflows.run_workflow("validate-groups-permissions")
99-
100-
101-
def test_permissions_migration_for_group_with_same_name(
102-
installation_ctx,
103-
make_cluster_policy,
104-
make_cluster_policy_permissions,
105-
):
106-
ws_group, acc_group = installation_ctx.make_ucx_group()
107-
migrated_group = MigratedGroup.partial_info(ws_group, acc_group)
108-
cluster_policy = make_cluster_policy()
109-
make_cluster_policy_permissions(
110-
object_id=cluster_policy.policy_id,
111-
permission_level=PermissionLevel.CAN_USE,
112-
group_name=migrated_group.name_in_workspace,
113-
)
114-
115-
schema_a = installation_ctx.make_schema()
116-
table_a = installation_ctx.make_table(schema_name=schema_a.name)
117-
installation_ctx.make_grant(migrated_group.name_in_workspace, 'USAGE', schema_info=schema_a)
118-
installation_ctx.make_grant(migrated_group.name_in_workspace, 'OWN', schema_info=schema_a)
119-
installation_ctx.make_grant(migrated_group.name_in_workspace, 'SELECT', table_info=table_a)
120-
121-
installation_ctx.workspace_installation.run()
122-
123-
installation_ctx.deployed_workflows.run_workflow("migrate-groups")
124-
125-
object_permissions = installation_ctx.generic_permissions_support.load_as_dict(
126-
"cluster-policies", cluster_policy.policy_id
127-
)
128-
new_schema_grants = installation_ctx.grants_crawler.for_schema_info(schema_a)
129-
130-
if {"USAGE", "OWN"} != new_schema_grants[migrated_group.name_in_account] or object_permissions[
131-
migrated_group.name_in_account
132-
] != PermissionLevel.CAN_USE:
133-
installation_ctx.deployed_workflows.relay_logs("migrate-groups")
134-
assert {"USAGE", "OWN"} == new_schema_grants[
135-
migrated_group.name_in_account
136-
], "Incorrect schema grants for migrated group"
137-
assert (
138-
object_permissions[migrated_group.name_in_account] == PermissionLevel.CAN_USE
139-
), "Incorrect permissions for migrated group"

0 commit comments

Comments
 (0)