|
| 1 | +import datetime as dt |
1 | 2 | import json |
2 | 3 | import logging |
3 | 4 | from collections import defaultdict |
| 5 | +from collections.abc import Callable, Iterable |
| 6 | +from functools import partial |
4 | 7 |
|
| 8 | +from databricks.sdk.retries import retried |
| 9 | + |
| 10 | +from databricks.labs.ucx.hive_metastore.grants import Grant |
5 | 11 | from databricks.labs.ucx.workspace_access.base import Permissions |
6 | 12 | from databricks.labs.ucx.workspace_access.groups import MigratedGroup, MigrationState |
7 | 13 | from databricks.labs.ucx.workspace_access.tacl import TableAclSupport |
@@ -35,29 +41,39 @@ def test_grants_with_permission_migration_api(runtime_ctx, ws, migrated_group, s |
35 | 41 | assert {"USAGE", "OWN"} == new_schema_grants["a"][migrated_group.name_in_account] |
36 | 42 |
|
37 | 43 |
|
38 | | -def test_permission_for_files_anonymous_func_migration_api(ws, sql_backend, runtime_ctx, migrated_group): |
39 | | - sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`") |
40 | | - sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`") |
| 44 | +def test_permission_for_files_anonymous_func_migration_api(runtime_ctx, migrated_group) -> None: |
| 45 | + """Test the migration of permissions for any files and anonymous functions.""" |
| 46 | + # TODO: Move migrated group into `runtime_ctx` and follow the `make_` pattern |
| 47 | + runtime_ctx.sql_backend.execute(f"GRANT READ_METADATA ON ANY FILE TO `{migrated_group.name_in_workspace}`") |
| 48 | + runtime_ctx.sql_backend.execute(f"GRANT SELECT ON ANONYMOUS FUNCTION TO `{migrated_group.name_in_workspace}`") |
41 | 49 |
|
42 | | - grants = runtime_ctx.grants_crawler |
| 50 | + MigrationState([migrated_group]).apply_to_groups_with_different_names(runtime_ctx.workspace_client) |
43 | 51 |
|
44 | | - MigrationState([migrated_group]).apply_to_groups_with_different_names(ws) |
| 52 | + @retried(on=[ValueError], timeout=dt.timedelta(minutes=2)) |
| 53 | + def verify_grants_meet_condition(condition: Callable[[Iterable[Grant]], bool], **kwargs) -> None: |
| 54 | + """Verify grants meet the condition. |
45 | 55 |
|
46 | | - any_file_actual = {} |
47 | | - for any_file_grant in grants.grants(any_file=True): |
48 | | - any_file_actual[any_file_grant.principal] = any_file_grant.action_type |
| 56 | + The method retries the condition check to account for eventual consistency of the permission API. |
| 57 | + """ |
| 58 | + grants = runtime_ctx.grants_crawler.grants(**kwargs) |
| 59 | + if not condition(grants): |
| 60 | + raise ValueError("Grants do not meet condition") |
49 | 61 |
|
50 | | - # both old and new group have permissions |
51 | | - assert migrated_group.name_in_workspace not in any_file_actual |
52 | | - assert migrated_group.name_in_account in any_file_actual |
| 62 | + def is_migrated_permission(grants: Iterable[Grant], *, permission: str) -> bool: |
| 63 | + """Check if the "any file" permissions are migrated. |
53 | 64 |
|
54 | | - anonymous_function_actual = {} |
55 | | - for ano_func_grant in grants.grants(anonymous_function=True): |
56 | | - anonymous_function_actual[ano_func_grant.principal] = ano_func_grant.action_type |
| 65 | + The migration permission APIs **move** the permission, thus the group name in workspace should not have the |
| 66 | + permissions anymore. |
| 67 | + """ |
| 68 | + principal_actions = {grant.principal: grant.action_type for grant in grants} |
| 69 | + return ( |
| 70 | + migrated_group.name_in_workspace not in principal_actions |
| 71 | + and migrated_group.name_in_account in principal_actions |
| 72 | + and principal_actions[migrated_group.name_in_account] == permission |
| 73 | + ) |
57 | 74 |
|
58 | | - assert migrated_group.name_in_workspace not in anonymous_function_actual |
59 | | - assert migrated_group.name_in_account in anonymous_function_actual |
60 | | - assert anonymous_function_actual[migrated_group.name_in_account] == "SELECT" |
| 75 | + verify_grants_meet_condition(partial(is_migrated_permission, permission="READ_METADATA"), any_file=True) |
| 76 | + verify_grants_meet_condition(partial(is_migrated_permission, permission="SELECT"), anonymous_function=True) |
61 | 77 |
|
62 | 78 |
|
63 | 79 | def test_permission_for_udfs_migration_api(ws, sql_backend, runtime_ctx, migrated_group): |
|
0 commit comments