Create a dashboard with a counter from a single query

src/databricks/labs/lsql/core.py

@@ -68,6 +68,7 @@ def as_dict(self) -> dict[str, Any]:
 
     # PySpark's compatibility
     def asDict(self, recursive: bool = False) -> dict[str, Any]:
+        _ = recursive
         return self.as_dict()
 
     def __eq__(self, other):

src/databricks/labs/lsql/dashboards.py → src/databricks/labs/lsql/dashboard.py

@@ -1,27 +1,47 @@
 import json
+import random
+import string
 from pathlib import Path
-from typing import Protocol, ClassVar, runtime_checkable
+from typing import ClassVar, Protocol, runtime_checkable
 
 import sqlglot
 import yaml
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.dashboards import Dashboard as SDKDashboard
 from databricks.sdk.service.workspace import ExportFormat
 
-from databricks.labs.lsql.lakeview import Dashboard, Page, Query, NamedQuery, ControlFieldEncoding
+from databricks.labs.lsql.lakeview import (
+    ControlFieldEncoding,
+    CounterEncodingMap,
+    CounterSpec,
+)
+from databricks.labs.lsql.lakeview import Dashboard as LakeviewDashboard
+from databricks.labs.lsql.lakeview import (
+    Dataset,
+    Field,
+    Layout,
+    NamedQuery,
+    Page,
+    Position,
+    Query,
+    Widget,
+)
+
 
 @runtime_checkable
 class _DataclassInstance(Protocol):
     __dataclass_fields__: ClassVar[dict]
 
-class Dashboards:
+
+class Dashboard:  # TODO: Rename, maybe DashboardClient?
     def __init__(self, ws: WorkspaceClient):
         self._ws = ws
 
     def get_dashboard(self, dashboard_path: str):
         with self._ws.workspace.download(dashboard_path, format=ExportFormat.SOURCE) as f:
-            raw = f.read().decode('utf-8')
+            raw = f.read().decode("utf-8")
             as_dict = json.loads(raw)
-            return Dashboard.from_dict(as_dict)
+            return LakeviewDashboard.from_dict(as_dict)
 
     def save_to_folder(self, dashboard_path: str, local_path: Path):
         local_path.mkdir(parents=True, exist_ok=True)
@@ -34,23 +54,67 @@ def save_to_folder(self, dashboard_path: str, local_path: Path):
             sql_query = dataset.query
             self._format_sql_file(sql_query, query_path)
         lvdash_yml = local_path / "lvdash.yml"
-        with lvdash_yml.open('w') as f:
+        with lvdash_yml.open("w") as f:
             first_page = dashboard.pages[0]
             self._replace_names(first_page, better_names)
             page = first_page.as_dict()
             yaml.safe_dump(page, f)
         assert True
 
+    @staticmethod
+    def _create_random_id() -> str:
+        charset = string.ascii_lowercase + string.digits
+        return "".join(random.choices(charset, k=8))
+
+    def create(self, dashboard_folder: Path) -> LakeviewDashboard:
+        """Create a dashboard from code, i.e. configuration and queries."""
+        datasets, layouts = [], []
+        for query_path in dashboard_folder.glob("*.sql"):
+            with query_path.open("r") as query_file:
+                raw_query = query_file.read()
+            dataset = Dataset(name=self._create_random_id(), display_name=query_path.stem, query=raw_query)
+            datasets.append(dataset)
+
+            fields = [Field(name="count", expression="`count`")]
+            query = Query(dataset_name=dataset.name, fields=fields)
+            named_query = NamedQuery(name=self._create_random_id(), query=query)
+            counter_spec = CounterSpec(CounterEncodingMap())
+            widget = Widget(name=self._create_random_id(), queries=[named_query], spec=counter_spec)
+            position = Position(x=0, y=0, width=1, height=1)
+            layout = Layout(widget=widget, position=position)
+            layouts.append(layout)
+
+        page = Page(name=dashboard_folder.name, display_name=dashboard_folder.name, layout=layouts)
+        lakeview_dashboard = LakeviewDashboard(datasets=datasets, pages=[page])
+        return lakeview_dashboard
+
+    def deploy(
+        self, lakeview_dashboard: LakeviewDashboard, *, display_name: str | None = None, dashboard_id: str | None = None
+    ) -> SDKDashboard:
+        """Deploy a lakeview dashboard."""
+        if (display_name is None and dashboard_id is None) or (display_name is not None and dashboard_id is not None):
+            raise ValueError("Give either display_name or dashboard_id.")
+        if display_name is not None:
+            dashboard = self._ws.lakeview.create(
+                display_name, serialized_dashboard=json.dumps(lakeview_dashboard.as_dict())
+            )
+        else:
+            assert dashboard_id is not None
+            dashboard = self._ws.lakeview.update(
+                dashboard_id, serialized_dashboard=json.dumps(lakeview_dashboard.as_dict())
+            )
+        return dashboard
+
     def _format_sql_file(self, sql_query, query_path):
-        with query_path.open('w') as f:
+        with query_path.open("w") as f:
             try:
                 for statement in sqlglot.parse(sql_query):
                     # see https://sqlglot.com/sqlglot/generator.html#Generator
                     pretty = statement.sql(
-                        dialect='databricks',
+                        dialect="databricks",
                         normalize=True,  # normalize identifiers to lowercase
                         pretty=True,  # format the produced SQL string
-                        normalize_functions='upper',  # normalize function names to uppercase
+                        normalize_functions="upper",  # normalize function names to uppercase
                         max_text_width=80,  # wrap text at 120 characters
                     )
                     f.write(f"{pretty};\n")
@@ -70,11 +134,11 @@ def _replace_names(self, node: _DataclassInstance, better_names: dict[str, str])
             node.dataset_name = better_names.get(node.dataset_name, node.dataset_name)
         elif isinstance(node, NamedQuery) and node.query:
             # 'dashboards/01eeb077e38c17e6ba3511036985960c/datasets/01eeb081882017f6a116991d124d3068_...'
-            if node.name.startswith('dashboards/'):
+            if node.name.startswith("dashboards/"):
                 parts = [node.query.dataset_name]
                 for field in node.query.fields:
                     parts.append(field.name)
-                new_name = '_'.join(parts)
+                new_name = "_".join(parts)
                 better_names[node.name] = new_name
             node.name = better_names.get(node.name, node.name)
         elif isinstance(node, ControlFieldEncoding):

src/databricks/labs/lsql/lakeview/model.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import abc
+import json
 from dataclasses import dataclass
 from enum import Enum
 from typing import Any
@@ -88,6 +89,10 @@ def as_dict(self) -> Json:
 
     @classmethod
     def from_dict(cls, d: Json) -> WidgetSpec:
+        if d["version"] == 0 and d["viz_spec"]["viz_type"].lower() == "table":
+            return TableV1Spec.from_dict(json.loads(d["viz_spec"]["serialized_options"]))
+        if d["version"] == 0 and d["viz_spec"]["viz_type"].lower() == "counter":
+            return CounterSpec.from_dict(json.loads(d["viz_spec"]["serialized_options"]))
         if d["version"] == 1 and d["widgetType"] == "details":
             return DetailsV1Spec.from_dict(d)
         if d["version"] == 1 and d["widgetType"] == "table":

tests/integration/queries/counter.sql

@@ -0,0 +1 @@
+SELECT 6217 AS count

tests/integration/test_dashboards.py

@@ -1,12 +1,57 @@
 from pathlib import Path
 
-from databricks.sdk import WorkspaceClient
+import pytest
 
-from databricks.labs.lsql.dashboards import Dashboards
+from databricks.labs.lsql.dashboard import Dashboard
+from databricks.labs.lsql.lakeview.model import CounterSpec
+
+
+@pytest.fixture
+def dashboard_id(ws, make_random):
+    """Clean the lakeview dashboard"""
+
+    dashboard_display_name = f"created_by_lsql_{make_random()}"
+    dashboard = ws.lakeview.create(dashboard_display_name)
+
+    yield dashboard.dashboard_id
+
+    ws.lakeview.trash(dashboard.dashboard_id)
 
 
 def test_load_dashboard(ws):
-    dashboards = Dashboards(ws)
+    dashboard = Dashboard(ws)
     src = "/Workspace/Users/[email protected]/Trivial Dashboard.lvdash.json"
     dst = Path(__file__).parent / "sample"
-    dashboards.save_to_folder(src, dst)
+    dashboard.save_to_folder(src, dst)
+
+
+def test_dashboard_deploys_one_dataset_per_query(ws):
+def test_dashboard_deploys_one_dataset_per_query(ws, make_random):
+def test_dashboard_creates_one_dataset_per_query(ws, make_random):
+def test_dashboard_creates_one_dataset_per_query(ws):
+    queries = Path(__file__).parent / "queries"
+    dashboard = Dashboard(ws).create(queries)
+    assert len(dashboard.datasets) == len([query for query in queries.glob("*.sql")])
+
+
+def test_dashboard_creates_one_counter_widget_per_query(ws):
+    queries = Path(__file__).parent / "queries"
+    dashboard = Dashboard(ws).create(queries)
+
+    counter_widgets = []
+    for page in dashboard.pages:
+        for layout in page.layout:
+            if isinstance(layout.widget.spec, CounterSpec):
+                counter_widgets.append(layout.widget)
+
+    assert len(counter_widgets) == len([query for query in queries.glob("*.sql")])
+
+
+def test_dashboard_deploys_dashboard(ws, dashboard_id):
+    queries = Path(__file__).parent / "queries"
+    dashboard_client = Dashboard(ws)
+    lakeview_dashboard = dashboard_client.create(queries)
+
+    dashboard = dashboard_client.deploy(lakeview_dashboard, dashboard_id=dashboard_id)
+
+    assert dashboard_client.get_dashboard(dashboard.path).as_dict() == lakeview_dashboard.as_dict()