diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py index 24ac1d61..e0038115 100644 --- a/src/databricks/labs/lsql/dashboards.py +++ b/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 diff --git a/tests/unit/test_dashboards.py b/tests/unit/test_dashboards.py index 97e1943d..bb773727 100644 --- a/tests/unit/test_dashboards.py +++ b/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)