Use widget width and height defined in query header (#147)

JCZuurmond · web-flow · commit 40e46e1fef31 · 2024-06-12T20:46:34.000+02:00
Resolves #114 <img width="1701" alt="Screenshot 2024-06-12 at 14 10 39" src="https://github.com/databrickslabs/lsql/assets/5946784/84f66335-bf42-4b3a-a52c-e0f3534d8fa2">
diff --git a/docs/dashboards.md b/docs/dashboards.md
@@ -79,8 +79,10 @@ when the metadata cannot be inferred from the query itself.
 
 ### Headers of SQL files
 
-The first line is used to define widget and vizualization metadata used to render the relevant portion of the dashboard.
-Metadata could be defined in `--` and `/* ... */` comments, which are detected by our SQL parser.
+The header is used to define widget and vizualization metadata used to render the relevant portion of the dashboard.
+Metadata could be defined in a `--` or `/* ... */` comment, which are detected by our SQL parser. The parser only reads
+the **comment starting on the top**, which can be a single line using `--` or span multiple lines
+using `/* ... */`.
 
 | Format       | Readability | Verbosity |
 |--------------|-------------|-----------|
@@ -89,6 +91,15 @@ Metadata could be defined in `--` and `/* ... */` comments, which are detected b
 | `argparse`   | ?           | lowest    |
 | Query string | ? | ? |
 
+#### Widget arguments
+
+The following widget arguments are supported:
+
+| Flag           | Description                                 | Type | Optional |
+|----------------|---------------------------------------------|------|----------|
+| -w or --width  | The number of columns that the widget spans | int  | Yes      |
+| -h or --height | The number of rows that the widget spans    | int  | Yes      |
+
 [[back to top](#dashboards-as-code)]
 
 ### Implicit detection
diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py
@@ -1,6 +1,9 @@
+import argparse
 import dataclasses
 import json
 import logging
+import shlex
+from argparse import ArgumentParser
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TypeVar
@@ -45,6 +48,35 @@ def as_dict(self) -> dict[str, str]:
         return dataclasses.asdict(self)
 
 
+@dataclass
+class WidgetMetadata:
+    width: int
+    height: int
+
+    def as_dict(self) -> dict[str, str]:
+        return dataclasses.asdict(self)
+
+    @staticmethod
+    def _get_arguments_parser() -> ArgumentParser:
+        parser = ArgumentParser("WidgetMetadata", add_help=False, exit_on_error=False)
+        parser.add_argument("-w", "--width", type=int)
+        parser.add_argument("-h", "--height", type=int)
+        return parser
+
+    def replace_from_arguments(self, arguments: list[str]) -> "WidgetMetadata":
+        parser = self._get_arguments_parser()
+        try:
+            args = parser.parse_args(arguments)
+        except (argparse.ArgumentError, SystemExit) as e:
+            logger.warning(f"Parsing {arguments}: {e}")
+            return dataclasses.replace(self)
+        return dataclasses.replace(
+            self,
+            width=args.width or self.width,
+            height=args.height or self.height,
+        )
+
+
 class Dashboards:
     _MAXIMUM_DASHBOARD_WIDTH = 6
 
@@ -119,7 +151,8 @@ def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
                     continue
             else:
                 widget = self._get_text_widget(path)
-            position = self._get_position(widget, position)
+            widget_metadata = self._parse_widget_metadata(path, widget)
+            position = self._get_position(widget_metadata, position)
             layout = Layout(widget=widget, position=position)
             layouts.append(layout)
 
@@ -150,6 +183,22 @@ def _parse_dashboard_metadata(dashboard_folder: Path) -> DashboardMetadata:
             logger.warning(f"Parsing {dashboard_metadata_path}: {e}")
             return fallback_metadata
 
+    def _parse_widget_metadata(self, path: Path, widget: Widget) -> WidgetMetadata:
+        width, height = self._get_width_and_height(widget)
+        fallback_metadata = WidgetMetadata(width, height)
+
+        try:
+            parsed_query = sqlglot.parse_one(path.read_text(), dialect=sqlglot.dialects.Databricks)
+        except sqlglot.ParseError as e:
+            logger.warning(f"Parsing {path}: {e}")
+            return fallback_metadata
+
+        if parsed_query.comments is None or len(parsed_query.comments) == 0:
+            return fallback_metadata
+
+        first_comment = parsed_query.comments[0]
+        return fallback_metadata.replace_from_arguments(shlex.split(first_comment))
+
     @staticmethod
     def _get_text_widget(path: Path) -> Widget:
         widget = Widget(name=path.stem, textbox_spec=path.read_text())
@@ -178,15 +227,14 @@ def _get_fields(query: str) -> list[Field]:
                 fields.append(field)
         return fields
 
-    def _get_position(self, widget: Widget, previous_position: Position) -> Position:
-        width, height = self._get_width_and_height(widget)
+    def _get_position(self, widget_metadata: WidgetMetadata, previous_position: Position) -> Position:
         x = previous_position.x + previous_position.width
-        if x + width > self._MAXIMUM_DASHBOARD_WIDTH:
+        if x + widget_metadata.width > self._MAXIMUM_DASHBOARD_WIDTH:
             x = 0
             y = previous_position.y + previous_position.height
         else:
             y = previous_position.y
-        position = Position(x=x, y=y, width=width, height=height)
+        position = Position(x=x, y=y, width=widget_metadata.width, height=widget_metadata.height)
         return position
 
     def _get_width_and_height(self, widget: Widget) -> tuple[int, int]:
diff --git a/tests/integration/test_dashboards.py b/tests/integration/test_dashboards.py
@@ -121,3 +121,17 @@ def test_dashboard_deploys_dashboard_with_counter_variation(ws, make_dashboard,
     sdk_dashboard = dashboards.deploy_dashboard(lakeview_dashboard, dashboard_id=sdk_dashboard.dashboard_id)
 
     assert ws.lakeview.get(sdk_dashboard.dashboard_id)
+
+
+def test_dashboard_deploys_dashboard_with_big_widget(ws, make_dashboard, tmp_path):
+    sdk_dashboard = make_dashboard()
+
+    query = """-- --width 6 --height 3\nSELECT 82917019218921 AS big_number_needs_big_widget"""
+    with (tmp_path / "counter.sql").open("w") as f:
+        f.write(query)
+    dashboards = Dashboards(ws)
+    lakeview_dashboard = dashboards.create_dashboard(tmp_path)
+
+    sdk_dashboard = dashboards.deploy_dashboard(lakeview_dashboard, dashboard_id=sdk_dashboard.dashboard_id)
+
+    assert ws.lakeview.get(sdk_dashboard.dashboard_id)
diff --git a/tests/unit/test_dashboards.py b/tests/unit/test_dashboards.py
@@ -1,11 +1,16 @@
+import dataclasses
 import logging
 from pathlib import Path
 from unittest.mock import create_autospec
 
 import pytest
 from databricks.sdk import WorkspaceClient
 
-from databricks.labs.lsql.dashboards import DashboardMetadata, Dashboards
+from databricks.labs.lsql.dashboards import (
+    DashboardMetadata,
+    Dashboards,
+    WidgetMetadata,
+)
 from databricks.labs.lsql.lakeview import (
     CounterEncodingMap,
     CounterSpec,
@@ -36,6 +41,29 @@ def test_dashboard_configuration_from_and_as_dict_is_a_unit_function():
     assert dashboard_metadata.as_dict() == raw
 
 
+def test_widget_metadata_replaces_arguments():
+    widget_metadata = WidgetMetadata(1, 1)
+    updated_metadata = widget_metadata.replace_from_arguments(["--width", "10", "--height", "10"])
+    assert updated_metadata.width == 10
+    assert updated_metadata.height == 10
+
+
+@pytest.mark.parametrize("attribute", ["width", "height"])
+def test_widget_metadata_replaces_one_attribute(attribute: str):
+    widget_metadata = WidgetMetadata(1, 1)
+    updated_metadata = widget_metadata.replace_from_arguments([f"--{attribute}", "10"])
+
+    other_fields = [field for field in dataclasses.fields(updated_metadata) if field.name != attribute]
+    assert getattr(updated_metadata, attribute) == 10
+    assert all(getattr(updated_metadata, field.name) == 1 for field in other_fields)
+
+
+def test_widget_metadata_as_dict():
+    raw = {"width": 10, "height": 10}
+    widget_metadata = WidgetMetadata(10, 10)
+    assert widget_metadata.as_dict() == raw
+
+
 def test_dashboards_saves_sql_files_to_folder(tmp_path):
     ws = create_autospec(WorkspaceClient)
     queries = Path(__file__).parent / "queries"
@@ -339,6 +367,94 @@ def test_dashboards_creates_dashboards_where_text_widget_has_expected_text(tmp_p
     ws.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    "header",
+    [
+        "-- --width 6 --height 3",
+        "/* --width 6 --height 3 */",
+        "/*\n--width 6\n--height 3 */",
+    ],
+)
+def test_dashboard_creates_dashboard_with_custom_sized_widget(tmp_path, header):
+    ws = create_autospec(WorkspaceClient)
+
+    query = f"{header}\nSELECT 82917019218921 AS big_number_needs_big_widget"
+    with (tmp_path / "counter.sql").open("w") as f:
+        f.write(query)
+
+    lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+    position = lakeview_dashboard.pages[0].layout[0].position
+
+    assert position.width == 6
+    assert position.height == 3
+    ws.assert_not_called()
+
+
+def test_dashboard_handles_incorrect_query_header(tmp_path, caplog):
+    ws = create_autospec(WorkspaceClient)
+
+    # Typo is on purpose
+    query = "-- --widh 6 --height 3 \nSELECT 82917019218921 AS big_number_needs_big_widget"
+    with (tmp_path / "counter.sql").open("w") as f:
+        f.write(query)
+
+    with caplog.at_level(logging.WARNING, logger="databricks.labs.lsql.dashboards"):
+        lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+
+    position = lakeview_dashboard.pages[0].layout[0].position
+    assert position.width == 1
+    assert position.height == 3
+    assert "--widh" in caplog.text
+    ws.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "query",
+    [
+        "-- --height 5\nSELECT 1 AS count -- --width 6",
+        "-- --height 5\nSELECT 1 AS count\n-- --width 6",
+        "-- --height 5\nSELECT 1 AS count\n/* --width 6 */",
+        "-- --height 5\n-- --width 6\nSELECT 1 AS count",
+        "-- --height 5\n/* --width 6 */\nSELECT 1 AS count",
+        "/* --height 5*/\n/* --width 6 */\nSELECT 1 AS count",
+        "/* --height 5*/\n-- --width 6 */\nSELECT 1 AS count",
+    ],
+)
+def test_dashboard_ignores_comment_on_other_lines(tmp_path, query):
+    ws = create_autospec(WorkspaceClient)
+
+    with (tmp_path / "counter.sql").open("w") as f:
+        f.write(query)
+
+    lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+    position = lakeview_dashboard.pages[0].layout[0].position
+
+    assert position.width == 1
+    assert position.height == 5
+    ws.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "query",
+    [
+        "SELECT 1\n-- --width 6 --height 6",
+        "SELECT 1\n/*\n--width 6\n--height 6*/",
+    ],
+)
+def test_dashboard_ignores_non_header_comment(tmp_path, query):
+    ws = create_autospec(WorkspaceClient)
+
+    with (tmp_path / "counter.sql").open("w") as f:
+        f.write(query)
+
+    lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+    position = lakeview_dashboard.pages[0].layout[0].position
+
+    assert position.width == 1
+    assert position.height == 3
+    ws.assert_not_called()
+
+
 def test_dashboards_deploy_calls_create_without_dashboard_id():
     ws = create_autospec(WorkspaceClient)
     dashboards = Dashboards(ws)
