[superset] 03/09: fix: avoid 500 errors with SQLLAB_BACKEND_PERSISTENCE (#25553)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 3.0
in repository https://gitbox.apache.org/repos/asf/superset.git

commit ec3bed709e5660cc2832826a194abdf3cacb3499
Author: Igor Khrol <kh...@gmail.com>
AuthorDate: Fri Oct 13 22:30:19 2023 +0300

    fix: avoid 500 errors with SQLLAB_BACKEND_PERSISTENCE (#25553)
    
    (cherry picked from commit 99f79f5143c417497ffde326a8393ab60aa71e7e)
---
 superset/views/sql_lab/views.py       |  4 ++++
 tests/integration_tests/core_tests.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/superset/views/sql_lab/views.py b/superset/views/sql_lab/views.py
index df7b36050e..08394eb2ba 100644
--- a/superset/views/sql_lab/views.py
+++ b/superset/views/sql_lab/views.py
@@ -219,6 +219,10 @@ class TabStateView(BaseSupersetView):
             return Response(status=403)
 
         fields = {k: json.loads(v) for k, v in request.form.to_dict().items()}
+        if client_id := fields.get("latest_query_id"):
+            query = db.session.query(Query).filter_by(client_id=client_id).one_or_none()
+            if not query:
+                return self.json_response({"error": "Bad request"}, status=400)
         db.session.query(TabState).filter_by(id=tab_state_id).update(fields)
         db.session.commit()
         return json_success(json.dumps(tab_state_id))
diff --git a/tests/integration_tests/core_tests.py b/tests/integration_tests/core_tests.py
index d9f998a066..191d9dc2d0 100644
--- a/tests/integration_tests/core_tests.py
+++ b/tests/integration_tests/core_tests.py
@@ -1197,6 +1197,41 @@ class TestCore(SupersetTestCase, InsertChartMixin):
 
         self.assertEqual(payload["label"], "Untitled Query foo")
 
+    def test_tabstate_update(self):
+        username = "admin"
+        self.login(username)
+        # create a tab
+        data = {
+            "queryEditor": json.dumps(
+                {
+                    "name": "Untitled Query foo",
+                    "dbId": 1,
+                    "schema": None,
+                    "autorun": False,
+                    "sql": "SELECT ...",
+                    "queryLimit": 1000,
+                }
+            )
+        }
+        resp = self.get_json_resp("/tabstateview/", data=data)
+        tab_state_id = resp["id"]
+        # update tab state with non-existing client_id
+        client_id = "asdfasdf"
+        data = {"sql": json.dumps("select 1"), "latest_query_id": json.dumps(client_id)}
+        response = self.client.put(f"/tabstateview/{tab_state_id}", data=data)
+        self.assertEqual(response.status_code, 400)
+        self.assertEqual(response.json["error"], "Bad request")
+        # generate query
+        db.session.add(Query(client_id=client_id, database_id=1))
+        db.session.commit()
+        # update tab state with a valid client_id
+        response = self.client.put(f"/tabstateview/{tab_state_id}", data=data)
+        self.assertEqual(response.status_code, 200)
+        # nulls should be ok too
+        data["latest_query_id"] = "null"
+        response = self.client.put(f"/tabstateview/{tab_state_id}", data=data)
+        self.assertEqual(response.status_code, 200)
+
     def test_virtual_table_explore_visibility(self):
         # test that default visibility it set to True
         database = superset.utils.database.get_example_database()