Skip to content

Commit

Permalink
Let migrate-dbsql-dashboards command to run as collection (#2656)
Browse files Browse the repository at this point in the history
## Changes
Let `migrate-dbsql-dashboards` command to run as collection

### Linked issues

Resolves #2612

### Functionality

- [x] modified existing command: `databricks labs ucx
migrate-dbsql-dashboards`

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] ~added integration tests~ : Covering after #2507
  • Loading branch information
JCZuurmond authored Sep 18, 2024
1 parent 6323815 commit fad0c3c
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 6 deletions.
2 changes: 2 additions & 0 deletions labs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,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.
Expand Down
16 changes: 13 additions & 3 deletions src/databricks/labs/ucx/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -563,10 +563,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
Expand Down
18 changes: 15 additions & 3 deletions tests/unit/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -799,9 +799,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):
Expand Down

0 comments on commit fad0c3c

Please sign in to comment.