This is an automated email from the ASF dual-hosted git repository. michaelsmolina pushed a commit to branch airbnb in repository https://gitbox.apache.org/repos/asf/superset.git
commit b1c7b819ae5fdad59aa50809d387a2d4c4778181 Author: Michael S. Molina <[email protected]> AuthorDate: Tue Nov 12 13:54:58 2024 -0300 fix: Exception handling for SQL Lab views (#30897) (cherry picked from commit c2885a166e961af62c0e15cd541133fc44b9669d) --- superset/views/sql_lab/views.py | 285 +++++++++++++++++++--------------- tests/integration_tests/core_tests.py | 1 - 2 files changed, 157 insertions(+), 129 deletions(-) diff --git a/superset/views/sql_lab/views.py b/superset/views/sql_lab/views.py index 29d0fdba75..29cf640d8c 100644 --- a/superset/views/sql_lab/views.py +++ b/superset/views/sql_lab/views.py @@ -29,11 +29,12 @@ from superset.constants import MODEL_VIEW_RW_METHOD_PERMISSION_MAP, RouteMethod from superset.models.sql_lab import Query, SavedQuery, TableSchema, TabState from superset.superset_typing import FlaskResponse from superset.utils import json -from superset.utils.core import get_user_id +from superset.utils.core import error_msg_from_exception, get_user_id from superset.views.base import ( BaseSupersetView, DeleteMixin, DeprecateModelViewMixin, + json_error_response, json_success, SupersetModelView, ) @@ -84,48 +85,56 @@ class TabStateView(BaseSupersetView): @has_access_api @expose("/", methods=("POST",)) def post(self) -> FlaskResponse: - query_editor = json.loads(request.form["queryEditor"]) - tab_state = TabState( - user_id=get_user_id(), - # This is for backward compatibility - label=query_editor.get("name") - or query_editor.get("title", __("Untitled Query")), - active=True, - database_id=query_editor["dbId"], - catalog=query_editor.get("catalog"), - schema=query_editor.get("schema"), - sql=query_editor.get("sql", "SELECT ..."), - query_limit=query_editor.get("queryLimit"), - hide_left_bar=query_editor.get("hideLeftBar"), - saved_query_id=query_editor.get("remoteId"), - template_params=query_editor.get("templateParams"), - ) - ( - db.session.query(TabState) - .filter_by(user_id=get_user_id()) - .update({"active": False}) - ) - db.session.add(tab_state) - db.session.commit() - return json_success(json.dumps({"id": tab_state.id})) + try: + query_editor = json.loads(request.form["queryEditor"]) + tab_state = TabState( + user_id=get_user_id(), + # This is for backward compatibility + label=query_editor.get("name") + or query_editor.get("title", __("Untitled Query")), + active=True, + database_id=query_editor["dbId"], + catalog=query_editor.get("catalog"), + schema=query_editor.get("schema"), + sql=query_editor.get("sql", "SELECT ..."), + query_limit=query_editor.get("queryLimit"), + hide_left_bar=query_editor.get("hideLeftBar"), + saved_query_id=query_editor.get("remoteId"), + template_params=query_editor.get("templateParams"), + ) + ( + db.session.query(TabState) + .filter_by(user_id=get_user_id()) + .update({"active": False}) + ) + db.session.add(tab_state) + db.session.commit() + return json_success(json.dumps({"id": tab_state.id})) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("/<int:tab_state_id>", methods=("DELETE",)) def delete(self, tab_state_id: int) -> FlaskResponse: - owner_id = _get_owner_id(tab_state_id) - if owner_id is None: - return Response(status=404) - if owner_id != get_user_id(): - return Response(status=403) - - db.session.query(TabState).filter(TabState.id == tab_state_id).delete( - synchronize_session=False - ) - db.session.query(TableSchema).filter( - TableSchema.tab_state_id == tab_state_id - ).delete(synchronize_session=False) - db.session.commit() - return json_success(json.dumps("OK")) + try: + owner_id = _get_owner_id(tab_state_id) + if owner_id is None: + return Response(status=404) + if owner_id != get_user_id(): + return Response(status=403) + + db.session.query(TabState).filter(TabState.id == tab_state_id).delete( + synchronize_session=False + ) + db.session.query(TableSchema).filter( + TableSchema.tab_state_id == tab_state_id + ).delete(synchronize_session=False) + db.session.commit() + return json_success(json.dumps("OK")) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("/<int:tab_state_id>", methods=("GET",)) @@ -146,19 +155,23 @@ class TabStateView(BaseSupersetView): @has_access_api @expose("<int:tab_state_id>/activate", methods=("POST",)) def activate(self, tab_state_id: int) -> FlaskResponse: - owner_id = _get_owner_id(tab_state_id) - if owner_id is None: - return Response(status=404) - if owner_id != get_user_id(): - return Response(status=403) - - ( - db.session.query(TabState) - .filter_by(user_id=get_user_id()) - .update({"active": TabState.id == tab_state_id}) - ) - db.session.commit() - return json_success(json.dumps(tab_state_id)) + try: + owner_id = _get_owner_id(tab_state_id) + if owner_id is None: + return Response(status=404) + if owner_id != get_user_id(): + return Response(status=403) + + ( + db.session.query(TabState) + .filter_by(user_id=get_user_id()) + .update({"active": TabState.id == tab_state_id}) + ) + db.session.commit() + return json_success(json.dumps(tab_state_id)) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("<int:tab_state_id>", methods=("PUT",)) @@ -169,102 +182,118 @@ class TabStateView(BaseSupersetView): if owner_id != get_user_id(): 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)) + try: + fields = {k: json.loads(v) for k, v in request.form.to_dict().items()} + db.session.query(TabState).filter_by(id=tab_state_id).update(fields) + db.session.commit() + return json_success(json.dumps(tab_state_id)) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("<int:tab_state_id>/migrate_query", methods=("POST",)) def migrate_query(self, tab_state_id: int) -> FlaskResponse: - owner_id = _get_owner_id(tab_state_id) - if owner_id is None: - return Response(status=404) - if owner_id != get_user_id(): - return Response(status=403) - - client_id = json.loads(request.form["queryId"]) - db.session.query(Query).filter_by(client_id=client_id).update( - {"sql_editor_id": tab_state_id} - ) - db.session.commit() - return json_success(json.dumps(tab_state_id)) + try: + owner_id = _get_owner_id(tab_state_id) + if owner_id is None: + return Response(status=404) + if owner_id != get_user_id(): + return Response(status=403) + + client_id = json.loads(request.form["queryId"]) + db.session.query(Query).filter_by(client_id=client_id).update( + {"sql_editor_id": tab_state_id} + ) + db.session.commit() + return json_success(json.dumps(tab_state_id)) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("<int:tab_state_id>/query/<client_id>", methods=("DELETE",)) def delete_query(self, tab_state_id: int, client_id: str) -> FlaskResponse: - # Before deleting the query, ensure it's not tied to any - # active tab as the last query. If so, replace the query - # with the latest one created in that tab - tab_state_query = db.session.query(TabState).filter_by( - id=tab_state_id, latest_query_id=client_id - ) - if tab_state_query.count(): - query = ( - db.session.query(Query) - .filter( - and_( - Query.client_id != client_id, - Query.user_id == get_user_id(), - Query.sql_editor_id == str(tab_state_id), - ), - ) - .order_by(Query.id.desc()) - .first() - ) - tab_state_query.update( - {"latest_query_id": query.client_id if query else None} + try: + # Before deleting the query, ensure it's not tied to any + # active tab as the last query. If so, replace the query + # with the latest one created in that tab + tab_state_query = db.session.query(TabState).filter_by( + id=tab_state_id, latest_query_id=client_id ) + if tab_state_query.count(): + query = ( + db.session.query(Query) + .filter( + and_( + Query.client_id != client_id, + Query.user_id == get_user_id(), + Query.sql_editor_id == str(tab_state_id), + ), + ) + .order_by(Query.id.desc()) + .first() + ) + tab_state_query.update( + {"latest_query_id": query.client_id if query else None} + ) - db.session.query(Query).filter_by( - client_id=client_id, - user_id=get_user_id(), - sql_editor_id=str(tab_state_id), - ).delete(synchronize_session=False) - db.session.commit() - return json_success(json.dumps("OK")) + db.session.query(Query).filter_by( + client_id=client_id, + user_id=get_user_id(), + sql_editor_id=str(tab_state_id), + ).delete(synchronize_session=False) + db.session.commit() + return json_success(json.dumps("OK")) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) class TableSchemaView(BaseSupersetView): @has_access_api @expose("/", methods=("POST",)) def post(self) -> FlaskResponse: - table = json.loads(request.form["table"]) - - # delete any existing table schema - db.session.query(TableSchema).filter( - TableSchema.tab_state_id == table["queryEditorId"], - TableSchema.database_id == table["dbId"], - TableSchema.catalog == table.get("catalog"), - TableSchema.schema == table["schema"], - TableSchema.table == table["name"], - ).delete(synchronize_session=False) - - table_schema = TableSchema( - tab_state_id=table["queryEditorId"], - database_id=table["dbId"], - catalog=table.get("catalog"), - schema=table["schema"], - table=table["name"], - description=json.dumps(table), - expanded=True, - ) - db.session.add(table_schema) - db.session.commit() - return json_success(json.dumps({"id": table_schema.id})) + try: + table = json.loads(request.form["table"]) + + # delete any existing table schema + db.session.query(TableSchema).filter( + TableSchema.tab_state_id == table["queryEditorId"], + TableSchema.database_id == table["dbId"], + TableSchema.catalog == table.get("catalog"), + TableSchema.schema == table["schema"], + TableSchema.table == table["name"], + ).delete(synchronize_session=False) + + table_schema = TableSchema( + tab_state_id=table["queryEditorId"], + database_id=table["dbId"], + catalog=table.get("catalog"), + schema=table["schema"], + table=table["name"], + description=json.dumps(table), + expanded=True, + ) + db.session.add(table_schema) + db.session.commit() + return json_success(json.dumps({"id": table_schema.id})) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("/<int:table_schema_id>", methods=("DELETE",)) def delete(self, table_schema_id: int) -> FlaskResponse: - db.session.query(TableSchema).filter(TableSchema.id == table_schema_id).delete( - synchronize_session=False - ) - db.session.commit() - return json_success(json.dumps("OK")) + try: + db.session.query(TableSchema).filter( + TableSchema.id == table_schema_id + ).delete(synchronize_session=False) + db.session.commit() + return json_success(json.dumps("OK")) + except Exception as ex: # pylint: disable=broad-except + db.session.rollback() + return json_error_response(error_msg_from_exception(ex), 400) @has_access_api @expose("/<int:table_schema_id>/expanded", methods=("POST",)) diff --git a/tests/integration_tests/core_tests.py b/tests/integration_tests/core_tests.py index 44b7ef26e6..2cc8c8f5ca 100644 --- a/tests/integration_tests/core_tests.py +++ b/tests/integration_tests/core_tests.py @@ -1015,7 +1015,6 @@ class TestCore(SupersetTestCase): 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()
