Use order key in query header if defined (#149)

Resolves #148 

<img width="1406" alt="Screenshot 2024-06-12 at 14 47 22"
src="https://github.com/databrickslabs/lsql/assets/5946784/43c44245-e7b7-46f4-9243-e138b1b18ef9">

src/databricks/labs/lsql/dashboards.py

@@ -4,6 +4,7 @@
 import logging
 import shlex
 from argparse import ArgumentParser
+from collections.abc import Iterable
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TypeVar
@@ -50,6 +51,7 @@ def as_dict(self) -> dict[str, str]:
 
 @dataclass
 class WidgetMetadata:
+    order: int
     width: int
     height: int
 
@@ -59,6 +61,7 @@ def as_dict(self) -> dict[str, str]:
     @staticmethod
     def _get_arguments_parser() -> ArgumentParser:
         parser = ArgumentParser("WidgetMetadata", add_help=False, exit_on_error=False)
+        parser.add_argument("-o", "--order", type=int)
         parser.add_argument("-w", "--width", type=int)
         parser.add_argument("-h", "--height", type=int)
         return parser
@@ -72,6 +75,7 @@ def replace_from_arguments(self, arguments: list[str]) -> "WidgetMetadata":
             return dataclasses.replace(self)
         return dataclasses.replace(
             self,
+            order=args.order or self.order,
             width=args.width or self.width,
             height=args.height or self.height,
         )
@@ -125,19 +129,32 @@ def _format_query(query: str) -> str:
         return formatted_query
 
     def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
-        """Create a dashboard from code, i.e. metadata and queries."""
+        """Create a dashboard from code, i.e. configuration and queries."""
         dashboard_metadata = self._parse_dashboard_metadata(dashboard_folder)
+        datasets = self._get_datasets(dashboard_folder)
+        widgets = self._get_widgets(dashboard_folder.iterdir(), datasets)
+        layouts = self._get_layouts(widgets)
+        page = Page(
+            name=dashboard_metadata.display_name,
+            display_name=dashboard_metadata.display_name,
+            layout=layouts,
+        )
+        lakeview_dashboard = Dashboard(datasets=datasets, pages=[page])
+        return lakeview_dashboard
 
-        position = Position(0, 0, 0, 0)  # First widget position
-        datasets, layouts = [], []
+    @staticmethod
+    def _get_datasets(dashboard_folder: Path) -> list[Dataset]:
+        datasets = []
         for query_path in sorted(dashboard_folder.glob("*.sql")):
             with query_path.open("r") as query_file:
                 raw_query = query_file.read()
             dataset = Dataset(name=query_path.stem, display_name=query_path.stem, query=raw_query)
             datasets.append(dataset)
+        return datasets
 
-        dataset_index = 0
-        for path in sorted(dashboard_folder.iterdir()):
+    def _get_widgets(self, files: Iterable[Path], datasets: list[Dataset]) -> list[tuple[Widget, WidgetMetadata]]:
+        dataset_index, widgets = 0, []
+        for order, path in enumerate(sorted(files)):
             if path.suffix not in {".sql", ".md"}:
                 continue
             if path.suffix == ".sql":
@@ -151,18 +168,17 @@ def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
                     continue
             else:
                 widget = self._get_text_widget(path)
-            widget_metadata = self._parse_widget_metadata(path, widget)
+            widget_metadata = self._parse_widget_metadata(path, widget, order)
+            widgets.append((widget, widget_metadata))
+        return widgets
+
+    def _get_layouts(self, widgets: list[tuple[Widget, WidgetMetadata]]) -> list[Layout]:
+        layouts, position = [], Position(0, 0, 0, 0)  # First widget position
+        for widget, widget_metadata in sorted(widgets, key=lambda w: (w[1].order, w[0].name)):
             position = self._get_position(widget_metadata, position)
             layout = Layout(widget=widget, position=position)
             layouts.append(layout)
-
-        page = Page(
-            name=dashboard_metadata.display_name,
-            display_name=dashboard_metadata.display_name,
-            layout=layouts,
-        )
-        lakeview_dashboard = Dashboard(datasets=datasets, pages=[page])
-        return lakeview_dashboard
+        return layouts
 
     @staticmethod
     def _parse_dashboard_metadata(dashboard_folder: Path) -> DashboardMetadata:
@@ -183,9 +199,13 @@ 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:
+    def _parse_widget_metadata(self, path: Path, widget: Widget, order: int) -> WidgetMetadata:
         width, height = self._get_width_and_height(widget)
-        fallback_metadata = WidgetMetadata(width, height)
+        fallback_metadata = WidgetMetadata(
+            order=order,
+            width=width,
+            height=height,
+        )
 
         try:
             parsed_query = sqlglot.parse_one(path.read_text(), dialect=sqlglot.dialects.Databricks)

tests/integration/test_dashboards.py

@@ -135,3 +135,23 @@ def test_dashboard_deploys_dashboard_with_big_widget(ws, make_dashboard, tmp_pat
     sdk_dashboard = dashboards.deploy_dashboard(lakeview_dashboard, dashboard_id=sdk_dashboard.dashboard_id)
 
     assert ws.lakeview.get(sdk_dashboard.dashboard_id)
+
+
+def test_dashboards_deploys_dashboard_with_order_overwrite(ws, make_dashboard, tmp_path):
+    sdk_dashboard = make_dashboard()
+
+    for query_name in range(6):
+        with (tmp_path / f"{query_name}.sql").open("w") as f:
+            f.write(f"SELECT {query_name} AS count")
+
+    # Move the '4' inbetween '1' and '2' query. Note that the order 1 puts '4' on the same position as '1', but with an
+    # order tiebreaker the query name decides the final order.
+    with (tmp_path / "4.sql").open("w") as f:
+        f.write("-- --order 1\nSELECT 4 AS count")
+
+    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)

tests/unit/test_dashboards.py

@@ -42,15 +42,15 @@ def test_dashboard_configuration_from_and_as_dict_is_a_unit_function():
 
 
 def test_widget_metadata_replaces_arguments():
-    widget_metadata = WidgetMetadata(1, 1)
+    widget_metadata = WidgetMetadata(1, 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"])
+@pytest.mark.parametrize("attribute", ["order", "width", "height"])
 def test_widget_metadata_replaces_one_attribute(attribute: str):
-    widget_metadata = WidgetMetadata(1, 1)
+    widget_metadata = WidgetMetadata(1, 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]
@@ -59,8 +59,8 @@ def test_widget_metadata_replaces_one_attribute(attribute: str):
 
 
 def test_widget_metadata_as_dict():
-    raw = {"width": 10, "height": 10}
-    widget_metadata = WidgetMetadata(10, 10)
+    raw = {"order": 10, "width": 10, "height": 10}
+    widget_metadata = WidgetMetadata(10, 10, 10)
     assert widget_metadata.as_dict() == raw
 
 
@@ -324,6 +324,25 @@ def test_dashboards_creates_dashboards_with_widgets_sorted_alphanumerically(tmp_
     ws.assert_not_called()
 
 
+def test_dashboards_creates_dashboards_with_widgets_order_overwrite(tmp_path):
+    ws = create_autospec(WorkspaceClient)
+
+    for query_name in "abcdf":
+        with (tmp_path / f"{query_name}.sql").open("w") as f:
+            f.write("SELECT 1 AS count")
+
+    # Move the 'e' inbetween 'b' and 'c' query. Note that the order 1 puts 'e' on the same position as 'b', but with an
+    # order tiebreaker the query name decides the final order.
+    with (tmp_path / "e.sql").open("w") as f:
+        f.write("-- --order 1\nSELECT 1 AS count")
+
+    lakeview_dashboard = Dashboards(ws).create_dashboard(tmp_path)
+    widget_names = [layout.widget.name for layout in lakeview_dashboard.pages[0].layout]
+
+    assert "".join(widget_names) == "abecdf"
+    ws.assert_not_called()
+
+
 @pytest.mark.parametrize("query, width, height", [("SELECT 1 AS count", 1, 3)])
 def test_dashboards_creates_dashboards_where_widget_has_expected_width_and_height(tmp_path, query, width, height):
     ws = create_autospec(WorkspaceClient)