Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Let migrate-acls command run as collection #2664

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion labs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -258,12 +258,14 @@ commands:

- name: migrate-acls
description: |
Migrate ACLs from legacy metastore to UC metastore.
Migrate access control lists from legacy metastore to UC metastore.
flags:
- name: target-catalog
description: (Optional) Target catalog to migrate ACLs to. Used for HMS-FED ACLs migration.
- name: hms-fed
description: (Optional) Migrate HMS-FED ACLs. If not provided, HMS ACLs will be migrated for migrated tables.
- name: run-as-collection
description: (Optional) Run the command for the collection of workspaces with ucx installed. Default is False.

- name: migrate-dbsql-dashboards
description: Migrate DBSQL dashboards by replacing legacy HMS tables in DBSQL queries with the corresponding new UC tables.
Expand Down
22 changes: 15 additions & 7 deletions src/databricks/labs/ucx/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -572,16 +572,24 @@ def migrate_tables(


@ucx.command
def migrate_acls(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None, **named_parameters):
def migrate_acls(
w: WorkspaceClient,
*,
ctx: WorkspaceContext | None = None,
run_as_collection: bool = False,
a: AccountClient | None = None,
**named_parameters,
):
"""
Migrate the ACLs for migrated tables and view. Can work with hms federation or other table migration scenarios.
"""
if ctx is None:
ctx = WorkspaceContext(w)
ctx.acl_migrator.migrate_acls(
target_catalog=named_parameters.get("target_catalog"),
hms_fed=named_parameters.get("hms_fed", False),
)
if ctx:
workspace_contexts = [ctx]
else:
workspace_contexts = _get_workspace_contexts(w, a, run_as_collection, **named_parameters)
target_catalog, hms_fed = named_parameters.get("target_catalog"), named_parameters.get("hms_fed", False)
for workspace_context in workspace_contexts:
workspace_context.acl_migrator.migrate_acls(target_catalog=target_catalog, hms_fed=hms_fed)


@ucx.command
Expand Down
28 changes: 26 additions & 2 deletions tests/unit/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
join_collection,
logs,
manual_workspace_info,
migrate_acls,
migrate_credentials,
migrate_dbsql_dashboards,
migrate_local_code,
Expand Down Expand Up @@ -100,6 +101,12 @@ def create_workspace_client_mock(workspace_id: int) -> WorkspaceClient:
}
}
),
'/Users/foo/.ucx/workspaces.json': json.dumps(
[
{'workspace_id': 123, 'workspace_name': '123'},
{'workspace_id': 456, 'workspace_name': '456'},
]
),
"/Users/foo/.ucx/uc_roles_access.csv": "role_arn,resource_type,privilege,resource_path\n"
"arn:aws:iam::123456789012:role/role_name,s3,READ_FILES,s3://labsawsbucket/",
"/Users/foo/.ucx/azure_storage_account_info.csv": "prefix,client_id,principal,privilege,type,directory_id\ntest,test,test,test,Application,test",
Expand Down Expand Up @@ -265,9 +272,9 @@ def test_manual_workspace_info(ws):
manual_workspace_info(ws, prompts)


def test_create_table_mapping(ws, acc_client):
def test_create_table_mapping_raises_value_error_because_no_tables_found(ws, acc_client) -> None:
ctx = WorkspaceContext(ws)
with pytest.raises(ValueError, match='databricks labs ucx sync-workspace-info'):
with pytest.raises(ValueError, match="No tables found. .*"):
create_table_mapping(ws, ctx, False, acc_client)


Expand Down Expand Up @@ -434,6 +441,23 @@ def test_save_storage_and_principal_gcp(ws):
principal_prefix_access(ws, ctx=ctx)


@pytest.mark.parametrize("run_as_collection", [True, False])
def test_migrate_acls_calls_workspace_id(
run_as_collection,
workspace_clients,
acc_client,
) -> None:
if not run_as_collection:
workspace_clients = [workspace_clients[0]]
migrate_acls(
workspace_clients[0],
run_as_collection=run_as_collection,
a=acc_client,
)
for workspace_client in workspace_clients:
workspace_client.get_workspace_id.assert_called()


def test_migrate_credentials_azure(ws, acc_client):
ws.config.is_azure = True
ws.workspace.upload.return_value = "test"
Expand Down