Add --spec flag (#188)

Resolves #187

src/databricks/labs/lsql/dashboards.py

@@ -9,6 +9,7 @@
 from argparse import ArgumentParser
 from collections.abc import Callable, Iterable, Sized
 from dataclasses import dataclass
+from enum import Enum, unique
 from pathlib import Path
 from typing import TypeVar
 
@@ -121,13 +122,23 @@ class QueryHandler(BaseHandler):
 
     @staticmethod
     def _get_arguments_parser() -> ArgumentParser:
-        parser = ArgumentParser("WidgetMetadata", add_help=False, exit_on_error=False)
+        parser = ArgumentParser("TileMetadata", add_help=False, exit_on_error=False)
         parser.add_argument("--id", type=str)
         parser.add_argument("-o", "--order", type=int)
         parser.add_argument("-w", "--width", type=int)
         parser.add_argument("-h", "--height", type=int)
         parser.add_argument("-t", "--title", type=str)
         parser.add_argument("-d", "--description", type=str)
+        parser.add_argument(
+            "-s",
+            "--spec",
+            type=lambda v: QuerySpec(v.upper()),
+            default=QuerySpec.AUTO,
+            help=(
+                "The widget spec to use, see classes with WidgetSpec as parent class in "
+                "databricks.labs.lsql.lakeview.model."
+            ),
+        )
         parser.add_argument(
             "-f",
             "--filter",
@@ -190,6 +201,24 @@ def split(self) -> tuple[str, str]:
         return "", self._content
 
 
+@unique
+class QuerySpec(str, Enum):
+    """The query widget spec"""
+
+    AUTO = "AUTO"
+    TABLE = "TABLE"
+    COUNTER = "COUNTER"
+
+    def as_widget_spec(self) -> type[WidgetSpec]:
+        widget_spec_mapping: dict[str, type[WidgetSpec]] = {
+            "TABLE": TableV2Spec,
+            "COUNTER": CounterSpec,
+        }
+        if self.name not in widget_spec_mapping:
+            raise ValueError(f"Can not convert to widget spec: {self}")
+        return widget_spec_mapping[self.name]
+
+
 class TileMetadata:
     def __init__(
         self,
@@ -200,6 +229,7 @@ def __init__(
         _id: str = "",
         title: str = "",
         description: str = "",
+        spec: QuerySpec = QuerySpec.AUTO,
         filters: list[str] | None = None,
     ):
         self._path = path
@@ -209,6 +239,7 @@ def __init__(
         self.id = _id or path.stem
         self.title = title
         self.description = description
+        self.spec = spec
         self.filters = filters or []
 
     def is_markdown(self) -> bool:
@@ -257,7 +288,7 @@ def from_path(cls, path: Path) -> "TileMetadata":
         return cls.from_dict(path=path, **header)
 
     def __repr__(self):
-        return f"WidgetMetdata<{self._path}>"
+        return f"TileMetadata<{self._path}>"
 
 
 class Tile:
@@ -384,6 +415,8 @@ def _find_fields(self) -> list[Field]:
 
     def infer_spec_type(self) -> type[WidgetSpec] | None:
         """Infer the spec type from the query."""
+        if self._tile_metadata.spec != QuerySpec.AUTO:
+            return self._tile_metadata.spec.as_widget_spec()
         fields = self._find_fields()
         if len(fields) == 0:
             return None

tests/unit/test_dashboards.py

@@ -16,6 +16,7 @@
     Dashboards,
     MarkdownHandler,
     QueryHandler,
+    QuerySpec,
     QueryTile,
     Tile,
     TileMetadata,
@@ -116,7 +117,8 @@ def test_query_handler_parses_empty_header(tmp_path):
 
     header = handler.parse_header()
 
-    assert all(value is None for value in header.values())
+    has_default = {"spec"}
+    assert all(value is None for key, value in header.items() if key not in has_default)
 
 
 @pytest.mark.parametrize(
@@ -156,7 +158,8 @@ def test_query_handler_ignores_non_header_comment(tmp_path, query):
 
     header = handler.parse_header()
 
-    assert all(value is None for value in header.values())
+    has_default = {"spec"}
+    assert all(value is None for key, value in header.items() if key not in has_default)
 
 
 @pytest.mark.parametrize("attribute", ["id", "order", "height", "width", "title", "description"])
@@ -170,6 +173,16 @@ def test_query_handler_parses_attribute_from_header(tmp_path, attribute):
     assert str(header[attribute]) == "10"
 
 
+def test_query_handler_parses_spec_attribute_from_header(tmp_path):
+    path = tmp_path / "query.sql"
+    path.write_text("-- --spec COUNTER\nSELECT 1")
+    handler = QueryHandler(path)
+
+    header = handler.parse_header()
+
+    assert header["spec"] == "COUNTER"
+
+
 @pytest.mark.parametrize(
     "query",
     [
@@ -250,6 +263,21 @@ def test_markdown_handler_warns_about_open_ended_header(tmp_path, caplog):
     assert content == body
 
 
+def test_query_spec_raises_value_error_when_converting_auto_to_widget_spec():
+    with pytest.raises(ValueError):
+        QuerySpec.AUTO.as_widget_spec()
+
+
+def test_query_spec_converts_all_to_widget_spec_except_auto():
+    for spec in QuerySpec:
+        if spec == QuerySpec.AUTO:
+            continue
+        try:
+            spec.as_widget_spec()
+        except ValueError as e:
+            assert False, e
+
+
 def test_tile_metadata_replaces_width_and_height(tmp_path):
     path = tmp_path / "test.sql"
     path.write_text("SELECT 1")
@@ -259,7 +287,7 @@ def test_tile_metadata_replaces_width_and_height(tmp_path):
     assert updated_metadata.height == 10
 
 
-@pytest.mark.parametrize("attribute", ["id", "order", "width", "height", "title", "description"])
+@pytest.mark.parametrize("attribute", ["id", "order", "width", "height", "title", "description", "spec"])
 def test_tile_metadata_replaces_attribute(tmp_path, attribute: str):
     path = tmp_path / "test.sql"
     path.write_text("SELECT 1")
@@ -271,6 +299,7 @@ def test_tile_metadata_replaces_attribute(tmp_path, attribute: str):
         _id="1",
         title="1",
         description="1",
+        spec=QuerySpec.AUTO,
     )
     updated_metadata = tile_metadata.from_dict(**{"path": path, attribute: "10"})
     assert str(getattr(updated_metadata, attribute)) == "10"
@@ -300,6 +329,7 @@ def test_tile_metadata_as_dict(tmp_path):
         "height": 6,
         "title": "Test widget",
         "description": "Longer explanation",
+        "spec": "auto",
         "filters": ["column"],
     }
     tile_metadata = TileMetadata(
@@ -309,6 +339,7 @@ def test_tile_metadata_as_dict(tmp_path):
         height=6,
         title="Test widget",
         description="Longer explanation",
+        spec="auto",
         filters=["column"],
     )
     assert tile_metadata.as_dict() == raw
@@ -643,9 +674,28 @@ def test_dashboards_creates_dashboard_with_expected_counter_field_encoding_names
     ws.assert_not_called()
 
 
+@pytest.mark.parametrize(
+    "query, spec_expected",
+    [
+        ("SELECT 1", CounterSpec),
+        ("SELECT 1, 2", TableV2Spec),
+        ("-- --spec auto\nSELECT 1, 2", TableV2Spec),
+        ("-- --spec counter\nSELECT 1, 2", CounterSpec),
+    ],
+)
+def test_dashboards_infers_query_spec(tmp_path, query, spec_expected):
+    (tmp_path / "query.sql").write_text(query)
+
+    ws = create_autospec(WorkspaceClient)
+    lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+
+    spec = lakeview_dashboard.pages[0].layout[0].widget.spec
+    assert isinstance(spec, spec_expected)
+    ws.assert_not_called()
+
+
 def test_dashboards_creates_dashboard_with_expected_table_field_encodings(tmp_path):
-    with (tmp_path / "query.sql").open("w") as f:
-        f.write("SELECT 1 AS first, 2 AS second")
+    (tmp_path / "query.sql").write_text("select 1 as first, 2 as second")
 
     ws = create_autospec(WorkspaceClient)
     lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
@@ -661,8 +711,7 @@ def test_dashboards_creates_dashboards_with_second_widget_to_the_right_of_the_fi
     ws = create_autospec(WorkspaceClient)
 
     for i in range(2):
-        with (tmp_path / f"counter_{i}.sql").open("w") as f:
-            f.write(f"SELECT {i} AS count")
+        (tmp_path / f"counter_{i}.sql").write_text(f"SELECT {i} AS count")
 
     lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)