From 5971fb55adc39a4b09b5328f60d7649b2776b317 Mon Sep 17 00:00:00 2001
From: Cor <jczuurmond@protonmail.com>
Date: Tue, 23 Jul 2024 11:17:52 +0200
Subject: [PATCH] Add publish flag to `Dashboards.create_dashboard` (#233)

Resolves #219
---
 labs.yml                               |  2 ++
 src/databricks/labs/lsql/cli.py        |  3 ++-
 src/databricks/labs/lsql/dashboards.py |  6 ++++++
 tests/unit/test_dashboards.py          | 14 ++++++++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/labs.yml b/labs.yml
index d83c5e9d..1ec4edd4 100644
--- a/labs.yml
+++ b/labs.yml
@@ -25,5 +25,7 @@ commands:
         description: |
           Overwrite the database in the queries' `FROM` clauses with given value. 
           Useful when developing with seperated databases, for example, for production and development.
+      - name: publish
+        description: Publish the dashboard after creating
       - name: no-open
         description: Do not open the dashboard in the browser after creating
diff --git a/src/databricks/labs/lsql/cli.py b/src/databricks/labs/lsql/cli.py
index 588fc451..271ab459 100644
--- a/src/databricks/labs/lsql/cli.py
+++ b/src/databricks/labs/lsql/cli.py
@@ -18,6 +18,7 @@ def create_dashboard(
     *,
     catalog: str = "",
     database: str = "",
+    publish: bool = False,
     no_open: bool = False,
 ):
     """Create a dashboard from queries"""
@@ -28,7 +29,7 @@ def create_dashboard(
         catalog=catalog or None,
         database=database or None,
     )
-    sdk_dashboard = lakeview_dashboards.create_dashboard(dashboard_metadata)
+    sdk_dashboard = lakeview_dashboards.create_dashboard(dashboard_metadata, publish=publish)
     if not no_open:
         assert sdk_dashboard.dashboard_id is not None
         dashboard_url = lakeview_dashboards.get_url(sdk_dashboard.dashboard_id)
diff --git a/src/databricks/labs/lsql/dashboards.py b/src/databricks/labs/lsql/dashboards.py
index 55ab5952..78dbcd65 100644
--- a/src/databricks/labs/lsql/dashboards.py
+++ b/src/databricks/labs/lsql/dashboards.py
@@ -901,6 +901,7 @@ def create_dashboard(
         parent_path: str | None = None,
         dashboard_id: str | None = None,
         warehouse_id: str | None = None,
+        publish: bool = False,
     ) -> SDKDashboard:
         """Create a Lakeview dashboard.
 
@@ -913,6 +914,8 @@ def create_dashboard(
                 The id of the dashboard to update
             warehouse_id : str | None (default: None)
                 The id of the warehouse to use
+            publish : bool (default: False)
+                If True, publish the dashboard. If False, save it as a draft.
         """
         dashboard_metadata.validate()
         serialized_dashboard = json.dumps(dashboard_metadata.as_lakeview().as_dict())
@@ -930,6 +933,9 @@ def create_dashboard(
                 serialized_dashboard=serialized_dashboard,
                 warehouse_id=warehouse_id,
             )
+        if publish:
+            assert sdk_dashboard.dashboard_id is not None
+            self._ws.lakeview.publish(sdk_dashboard.dashboard_id, warehouse_id=warehouse_id)
         return sdk_dashboard
 
     def deploy_dashboard(self, dashboard: Dashboard, **kwargs) -> SDKDashboard:
diff --git a/tests/unit/test_dashboards.py b/tests/unit/test_dashboards.py
index 27bbeb66..62d02c12 100644
--- a/tests/unit/test_dashboards.py
+++ b/tests/unit/test_dashboards.py
@@ -9,6 +9,7 @@
 import sqlglot
 import yaml
 from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.dashboards import Dashboard as SDKDashboard
 
 from databricks.labs.lsql.dashboards import (
     BaseHandler,
@@ -1246,6 +1247,7 @@ def test_dashboards_calls_create_without_dashboard_id():
         warehouse_id="warehouse",
     )
     ws.lakeview.update.assert_not_called()
+    ws.lakeview.publish.assert_not_called()
 
 
 def test_dashboards_calls_update_with_dashboard_id():
@@ -1262,6 +1264,18 @@ def test_dashboards_calls_update_with_dashboard_id():
         serialized_dashboard=json.dumps({"pages": [{"displayName": "test", "name": "test"}]}),
         warehouse_id="warehouse",
     )
+    ws.lakeview.publish.assert_not_called()
+
+
+def test_dashboards_calls_publish_with_dashboard_id():
+    ws = create_autospec(WorkspaceClient)
+    ws.lakeview.create.return_value = SDKDashboard(dashboard_id="id")
+    dashboards = Dashboards(ws)
+    dashboard_metadata = DashboardMetadata("test")
+
+    dashboards.create_dashboard(dashboard_metadata, warehouse_id="warehouse", publish=True)
+
+    ws.lakeview.publish.assert_called_with("id", warehouse_id="warehouse")
 
 
 def test_dashboard_raises_value_error_when_creating_dashboard_with_invalid_queries(tmp_path):