fix(sqllab): invalid persisted tab state (apache#25308)

10webio · Sep 19, 2023 · 7aaae06 · 7aaae06
1 parent b722a95
commit 7aaae06
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 13 deletions.
diff --git a/superset/sqllab/utils.py b/superset/sqllab/utils.py
@@ -79,19 +79,8 @@ def write_ipc_buffer(table: pa.Table) -> pa.Buffer:
 
 
 def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]:
-    # send list of tab state ids
-    tabs_state = (
-        db.session.query(TabState.id, TabState.label).filter_by(user_id=user_id).all()
-    )
-    tab_state_ids = [str(tab_state[0]) for tab_state in tabs_state]
-    # return first active tab, or fallback to another one if no tab is active
-    active_tab = (
-        db.session.query(TabState)
-        .filter_by(user_id=user_id)
-        .order_by(TabState.active.desc())
-        .first()
-    )
-
+    tabs_state: list[Any] = []
+    active_tab: Any = None
     databases: dict[int, Any] = {}
     for database in DatabaseDAO.find_all():
         databases[database.id] = {
@@ -102,6 +91,20 @@ def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]:
 
     # These are unnecessary if sqllab backend persistence is disabled
     if is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE"):
+        # send list of tab state ids
+        tabs_state = (
+            db.session.query(TabState.id, TabState.label)
+            .filter_by(user_id=user_id)
+            .all()
+        )
+        tab_state_ids = [str(tab_state[0]) for tab_state in tabs_state]
+        # return first active tab, or fallback to another one if no tab is active
+        active_tab = (
+            db.session.query(TabState)
+            .filter_by(user_id=user_id)
+            .order_by(TabState.active.desc())
+            .first()
+        )
         # return all user queries associated with existing SQL editors
         user_queries = (
             db.session.query(Query)

diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py
@@ -61,6 +61,35 @@ def test_get_from_empty_bootsrap_data(self):
         assert result["tab_state_ids"] == []
         self.assertEqual(len(result["databases"]), 0)
 
+    @mock.patch.dict(
+        "superset.extensions.feature_flag_manager._feature_flags",
+        {"SQLLAB_BACKEND_PERSISTENCE": False},
+        clear=True,
+    )
+    def test_get_from_bootstrap_data_for_non_persisted_tab_state(self):
+        self.login("admin")
+        # create a tab
+        data = {
+            "queryEditor": json.dumps(
+                {
+                    "title": "Untitled Query 1",
+                    "dbId": 1,
+                    "schema": None,
+                    "autorun": False,
+                    "sql": "SELECT ...",
+                    "queryLimit": 1000,
+                }
+            )
+        }
+        self.get_json_resp("/tabstateview/", data=data)
+        resp = self.client.get("/api/v1/sqllab/")
+        assert resp.status_code == 200
+        data = json.loads(resp.data.decode("utf-8"))
+        result = data.get("result")
+        assert result["active_tab"] == None
+        assert result["queries"] == {}
+        assert result["tab_state_ids"] == []
+
     @mock.patch.dict(
         "superset.extensions.feature_flag_manager._feature_flags",
         {"SQLLAB_BACKEND_PERSISTENCE": True},