Programmatically create a dashboard (#121)

This PR adds ability to programmatically create a dashboard.
Vizualization and widget options are reverse-engineered from POST
statements sent through UI.

<img width="1844" alt="image"
src="https://github.com/databricks/ucx/assets/259697/7c106e08-c49b-4abb-a36c-dc7dcd7fa713">

Author: nfx

pyproject.toml

@@ -176,6 +176,7 @@ ban-relative-imports = "all"
 [tool.ruff.per-file-ignores]
 # Tests can use magic values, assertions, and relative imports
 "tests/**/*" = ["PLR2004", "S101", "TID252"]
+"src/databricks/labs/ucx/providers/mixins/redash.py" = ["A002", "A003", "N815"]
 
 [tool.coverage.run]
 branch = true

src/databricks/labs/ucx/providers/mixins/compute.py

@@ -70,11 +70,12 @@ def __init__(self, ws: WorkspaceClient, cluster_id: str | None = None, language:
     def run(self, code):
         code = self._trim_leading_whitespace(code)
 
-        # perform AST transformations for very repetitive tasks, like JSON serialization
-        code_tree = ast.parse(code)
-        json_serialize_transform = _ReturnToPrintJsonTransformer()
-        new_tree = json_serialize_transform.apply(code_tree)
-        code = ast.unparse(new_tree)
+        if self._language == Language.PYTHON:
+            # perform AST transformations for very repetitive tasks, like JSON serialization
+            code_tree = ast.parse(code)
+            json_serialize_transform = _ReturnToPrintJsonTransformer()
+            new_tree = json_serialize_transform.apply(code_tree)
+            code = ast.unparse(new_tree)
 
         ctx = self._running_command_context()
         result = self._commands.execute(
@@ -84,7 +85,11 @@ def run(self, code):
         results = result.results
         if result.status == compute.CommandStatus.FINISHED:
             self._raise_if_failed(results)
-            if results.result_type == compute.ResultType.TEXT and json_serialize_transform.has_return:
+            if (
+                self._language == Language.PYTHON
+                and results.result_type == compute.ResultType.TEXT
+                and json_serialize_transform.has_return
+            ):
                 # parse json from converted return statement
                 return json.loads(results.data)
             return results.data
@@ -95,9 +100,9 @@ def run(self, code):
     def install_notebook_library(self, library):
         return self.run(
             f"""
-        get_ipython().run_line_magic('pip', 'install {library}')
-        dbutils.library.restartPython()
-        """
+            get_ipython().run_line_magic('pip', 'install {library}')
+            dbutils.library.restartPython()
+            """
         )
 
     def _running_command_context(self) -> compute.ContextStatusResponse:

src/databricks/labs/ucx/providers/mixins/redash.py

@@ -0,0 +1,272 @@
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from databricks.sdk.service._internal import _from_dict
+from databricks.sdk.service.sql import Visualization, Widget
+
+
+@dataclass
+class WidgetOptions:
+    created_at: str | None = None
+    description: str | None = None
+    is_hidden: bool | None = None
+    parameter_mappings: Any | None = None
+    position: Optional["WidgetPosition"] = None
+    title: str | None = None
+    updated_at: str | None = None
+
+    def as_dict(self) -> dict:
+        body = {}
+        if self.created_at is not None:
+            body["created_at"] = self.created_at
+        if self.description is not None:
+            body["description"] = self.description
+        if self.is_hidden is not None:
+            body["isHidden"] = self.is_hidden
+        if self.parameter_mappings:
+            body["parameterMappings"] = self.parameter_mappings
+        if self.position:
+            body["position"] = self.position.as_dict()
+        if self.title is not None:
+            body["title"] = self.title
+        if self.updated_at is not None:
+            body["updated_at"] = self.updated_at
+        return body
+
+    @classmethod
+    def from_dict(cls, d: dict[str, any]) -> "WidgetOptions":
+        return cls(
+            created_at=d.get("created_at", None),
+            description=d.get("description", None),
+            is_hidden=d.get("isHidden", None),
+            parameter_mappings=d.get("parameterMappings", None),
+            position=_from_dict(d, "position", WidgetPosition),
+            title=d.get("title", None),
+            updated_at=d.get("updated_at", None),
+        )
+
+
+@dataclass
+class WidgetPosition:
+    """Coordinates of this widget on a dashboard. This portion of the API changes frequently and is
+    unsupported."""
+
+    auto_height: bool | None = None
+    col: int | None = None
+    row: int | None = None
+    size_x: int | None = None
+    size_y: int | None = None
+
+    def as_dict(self) -> dict:
+        body = {}
+        if self.auto_height is not None:
+            body["autoHeight"] = self.auto_height
+        if self.col is not None:
+            body["col"] = self.col
+        if self.row is not None:
+            body["row"] = self.row
+        if self.size_x is not None:
+            body["sizeX"] = self.size_x
+        if self.size_y is not None:
+            body["sizeY"] = self.size_y
+        return body
+
+    @classmethod
+    def from_dict(cls, d: dict[str, any]) -> "WidgetPosition":
+        return cls(
+            auto_height=d.get("autoHeight", None),
+            col=d.get("col", None),
+            row=d.get("row", None),
+            size_x=d.get("sizeX", None),
+            size_y=d.get("sizeY", None),
+        )
+
+
+class DashboardWidgetsAPI:
+    """This is an evolving API that facilitates the addition and removal of widgets from existing dashboards
+    within the Databricks Workspace. Data structures may change over time."""
+
+    def __init__(self, api_client):
+        self._api = api_client
+
+    def create(
+        self,
+        dashboard_id: str,
+        options: WidgetOptions,
+        *,
+        text: str | None = None,
+        visualization_id: str | None = None,
+        width: int | None = None,
+    ) -> Widget:
+        """Add widget to a dashboard.
+
+        :param dashboard_id: str
+          Dashboard ID returned by :method:dashboards/create.
+        :param options: :class:`WidgetOptions` (optional)
+        :param text: str (optional)
+          If this is a textbox widget, the application displays this text. This field is ignored if the widget
+          contains a visualization in the `visualization` field.
+        :param visualization_id: str (optional)
+          Query Vizualization ID returned by :method:queryvisualizations/create.
+        :param width: int (optional)
+          Width of a widget
+
+        :returns: :class:`Widget`
+        """
+        body = {}
+        if dashboard_id is not None:
+            body["dashboard_id"] = dashboard_id
+        if options is not None:
+            body["options"] = options.as_dict()
+        if text is not None:
+            body["text"] = text
+        if visualization_id is not None:
+            body["visualization_id"] = visualization_id
+        if width is not None:
+            body["width"] = width
+        res = self._api.do("POST", "/api/2.0/preview/sql/widgets", body=body)
+        return Widget.from_dict(res)
+
+    def delete(self, id: str):
+        self._api.do("DELETE", f"/api/2.0/preview/sql/widgets/{id}")
+
+    def update(
+        self,
+        dashboard_id: str,
+        id: str,
+        *,
+        options: WidgetOptions | None = None,
+        text: str | None = None,
+        visualization_id: str | None = None,
+        width: int | None = None,
+    ) -> Widget:
+        """Update existing widget.
+
+        :param dashboard_id: str
+          Dashboard ID returned by :method:dashboards/create.
+        :param id: str
+        :param options: :class:`WidgetOptions` (optional)
+        :param text: str (optional)
+          If this is a textbox widget, the application displays this text. This field is ignored if the widget
+          contains a visualization in the `visualization` field.
+        :param visualization_id: str (optional)
+          Query Vizualization ID returned by :method:queryvisualizations/create.
+        :param width: int (optional)
+          Width of a widget
+
+        :returns: :class:`Widget`
+        """
+        body = {}
+        if dashboard_id is not None:
+            body["dashboard_id"] = dashboard_id
+        if options is not None:
+            body["options"] = options.as_dict()
+        if text is not None:
+            body["text"] = text
+        if visualization_id is not None:
+            body["visualization_id"] = visualization_id
+        if width is not None:
+            body["width"] = width
+        res = self._api.do("POST", f"/api/2.0/preview/sql/widgets/{id}", body=body)
+        return Widget.from_dict(res)
+
+
+class QueryVisualizationsAPI:
+    """This is an evolving API that facilitates the addition and removal of vizualisations from existing queries
+    within the Databricks Workspace. Data structures may change over time."""
+
+    def __init__(self, api_client):
+        self._api = api_client
+
+    def create(
+        self,
+        query_id: str,
+        type: str,
+        options: dict,
+        *,
+        created_at: str | None = None,
+        description: str | None = None,
+        name: str | None = None,
+        updated_at: str | None = None,
+    ) -> Visualization:
+        body = {}
+        if query_id is not None:
+            body["query_id"] = query_id
+        if type is not None:
+            body["type"] = type
+        if options is not None:
+            body["options"] = options
+        if name is not None:
+            body["name"] = name
+        if created_at is not None:
+            body["created_at"] = created_at
+        if description is not None:
+            body["description"] = description
+        if updated_at is not None:
+            body["updated_at"] = updated_at
+        res = self._api.do("POST", "/api/2.0/preview/sql/visualizations", body=body)
+        return Visualization.from_dict(res)
+
+    def delete(self, id: str):
+        """Remove visualization.
+
+        :param id: str
+        """
+
+        headers = {
+            "Accept": "application/json",
+        }
+        self._api.do("DELETE", f"/api/2.0/preview/sql/visualizations/{id}", headers=headers)
+
+
+@dataclass
+class VizColumn:
+    name: str
+    title: str
+    type: str = "string"
+    imageUrlTemplate: str = "{{ @ }}"
+    imageTitleTemplate: str = "{{ @ }}"
+    linkUrlTemplate: str = "{{ @ }}"
+    linkTextTemplate: str = "{{ @ }}"
+    linkTitleTemplate: str = "{{ @ }}"
+    linkOpenInNewTab: bool = True
+    displayAs: str = "string"
+    visible: bool = True
+    order: int = 100000
+    allowSearch: bool = False
+    alignContent: str = "left"
+    allowHTML: bool = False
+    highlightLinks: bool = False
+    useMonospaceFont: bool = False
+    preserveWhitespace: bool = False
+
+    def as_dict(self):
+        return dataclasses.asdict(self)
+
+
+class QueryVisualizationsExt(QueryVisualizationsAPI):
+    def create_table(
+        self,
+        query_id: str,
+        name: str,
+        columns: list[VizColumn],
+        *,
+        items_per_page: int = 25,
+        condensed=True,
+        with_row_number=False,
+        description: str | None = None,
+    ):
+        return self.create(
+            query_id,
+            "TABLE",
+            {
+                "itemsPerPage": items_per_page,
+                "condensed": condensed,
+                "withRowNumber": with_row_number,
+                "version": 2,
+                "columns": [x.as_dict() for x in columns],
+            },
+            name=name,
+            description=description,
+        )

tests/integration/test_dashboards.py

@@ -0,0 +1,60 @@
+import os
+
+import pytest
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.sql import AccessControl, ObjectTypePlural, PermissionLevel
+
+from databricks.labs.ucx.providers.mixins.redash import (
+    DashboardWidgetsAPI,
+    QueryVisualizationsExt,
+    VizColumn,
+    WidgetOptions,
+    WidgetPosition,
+)
+
+# logging.getLogger("databricks").setLevel("DEBUG")
+
+
+def test_creating_widgets(ws: WorkspaceClient):
+    pytest.skip()
+    dashboard_widgets_api = DashboardWidgetsAPI(ws.api_client)
+    query_visualizations_api = QueryVisualizationsExt(ws.api_client)
+
+    x = ws.dashboards.create(name="test dashboard")
+    ws.dbsql_permissions.set(
+        ObjectTypePlural.DASHBOARDS,
+        x.id,
+        access_control_list=[AccessControl(group_name="users", permission_level=PermissionLevel.CAN_MANAGE)],
+    )
+
+    dashboard_widgets_api.create(
+        x.id,
+        WidgetOptions(
+            title="first widget",
+            description="description of the widget",
+            position=WidgetPosition(col=0, row=0, size_x=3, size_y=3),
+        ),
+        text="this is _some_ **markdown**",
+        width=1,
+    )
+
+    dashboard_widgets_api.create(
+        x.id,
+        WidgetOptions(title="second", position=WidgetPosition(col=0, row=3, size_x=3, size_y=3)),
+        text="another text",
+        width=1,
+    )
+
+    data_sources = {x.warehouse_id: x.id for x in ws.data_sources.list()}
+    warehouse_id = os.environ["TEST_DEFAULT_WAREHOUSE_ID"]
+
+    query = ws.queries.create(
+        data_source_id=data_sources[warehouse_id],
+        description="abc",
+        name="this is a test query",
+        query="SHOW DATABASES",
+        run_as_role="viewer",
+    )
+
+    y = query_visualizations_api.create_table(query.id, "ABC Viz", [VizColumn(name="databaseName", title="DB")])
+    print(y)