From 07311b6d18b5888a17558b0d1bc51b91dc4a59ce Mon Sep 17 00:00:00 2001
From: Cor Zuurmond <jczuurmond@protonmail.com>
Date: Wed, 12 Jun 2024 14:37:10 +0200
Subject: [PATCH] Implement widget ordering

---
 src/databricks/labs/lsql/dashboards.py | 22 ++++++++++++++++------
 tests/unit/test_dashboards.py          | 19 +++++++++++++++++++
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py
index ec5d14f9..3fdd2839 100644
--- a/src/databricks/labs/lsql/dashboards.py
+++ b/src/databricks/labs/lsql/dashboards.py
@@ -36,12 +36,14 @@
 
 @dataclass
 class WidgetMetadata:
+    order: int
     width: int
     height: int
 
     @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
@@ -55,6 +57,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,
         )
@@ -109,7 +112,6 @@ def _format_query(query: str) -> str:
 
     def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
         """Create a dashboard from code, i.e. configuration and queries."""
-        position = Position(0, 0, 0, 0)  # First widget position
         datasets, layouts = [], []
         for query_path in sorted(dashboard_folder.glob("*.sql")):
             with query_path.open("r") as query_file:
@@ -117,8 +119,8 @@ def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
             dataset = Dataset(name=query_path.stem, display_name=query_path.stem, query=raw_query)
             datasets.append(dataset)
 
-        dataset_index = 0
-        for path in sorted(dashboard_folder.iterdir()):
+        dataset_index, widgets = 0, []
+        for order, path in enumerate(sorted(dashboard_folder.iterdir())):
             if path.suffix not in {".sql", ".md"}:
                 continue
             if path.suffix == ".sql":
@@ -132,7 +134,11 @@ 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))
+
+        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)
@@ -141,9 +147,13 @@ def create_dashboard(self, dashboard_folder: Path) -> Dashboard:
         lakeview_dashboard = Dashboard(datasets=datasets, pages=[page])
         return lakeview_dashboard
 
-    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)
diff --git a/tests/unit/test_dashboards.py b/tests/unit/test_dashboards.py
index 370e74b7..f4127013 100644
--- a/tests/unit/test_dashboards.py
+++ b/tests/unit/test_dashboards.py
@@ -262,6 +262,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)