diff --git a/labs.yml b/labs.yml index 2ee694da06..5d64fc84a8 100644 --- a/labs.yml +++ b/labs.yml @@ -262,6 +262,8 @@ commands: flags: - name: dashboard-id description: (Optional) DBSQL dashboard ID to migrate. If no dashboard ID is provided, all DBSQL dashboards in the workspace will be migrated. + - name: run-as-collection + description: (Optional) Run the command for the collection of workspaces with ucx installed. Default is False. - name: revert-dbsql-dashboards description: Revert DBSQL dashboards that have been migrated to their original state before the migration. diff --git a/src/databricks/labs/ucx/cli.py b/src/databricks/labs/ucx/cli.py index e5b97a4ceb..6e2a1a6875 100644 --- a/src/databricks/labs/ucx/cli.py +++ b/src/databricks/labs/ucx/cli.py @@ -554,10 +554,20 @@ def migrate_acls(w: WorkspaceClient, *, ctx: WorkspaceContext | None = None, **n @ucx.command -def migrate_dbsql_dashboards(w: WorkspaceClient, dashboard_id: str | None = None): +def migrate_dbsql_dashboards( + w: WorkspaceClient, + dashboard_id: str | None = None, + ctx: WorkspaceContext | None = None, + run_as_collection: bool = False, + a: AccountClient | None = None, +) -> None: """Migrate table references in DBSQL Dashboard queries""" - ctx = WorkspaceContext(w) - ctx.redash.migrate_dashboards(dashboard_id) + if ctx: + workspace_contexts = [ctx] + else: + workspace_contexts = _get_workspace_contexts(w, a, run_as_collection) + for workspace_context in workspace_contexts: + workspace_context.redash.migrate_dashboards(dashboard_id) @ucx.command diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 3a687eb692..30f5212279 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -787,9 +787,21 @@ def test_create_missing_principal_azure(ws, caplog): assert str(failure.value) == "Unsupported cloud provider" -def test_migrate_dbsql_dashboards(ws, caplog): - migrate_dbsql_dashboards(ws) - ws.dashboards.list.assert_called_once() +@pytest.mark.parametrize("run_as_collection", [False, True]) +def test_migrate_dbsql_dashboards_list_dashboards( + run_as_collection, + workspace_clients, + acc_client, +) -> None: + if not run_as_collection: + workspace_clients = [workspace_clients[0]] + migrate_dbsql_dashboards( + workspace_clients[0], + run_as_collection=run_as_collection, + a=acc_client, + ) + for workspace_client in workspace_clients: + workspace_client.dashboards.list.assert_called_once() def test_revert_dbsql_dashboards(ws, caplog):