fix: Exception handling for SQL Lab views (#30897)

This commit is contained in:
Michael S. Molina 2024-11-12 13:54:58 -03:00 committed by GitHub
parent 0e165c1a21
commit c2885a166e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 152 additions and 124 deletions

View File

@ -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,6 +85,7 @@ class TabStateView(BaseSupersetView):
@has_access_api
@expose("/", methods=("POST",))
def post(self) -> FlaskResponse:
try:
query_editor = json.loads(request.form["queryEditor"])
tab_state = TabState(
user_id=get_user_id(),
@ -108,10 +110,14 @@ class TabStateView(BaseSupersetView):
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:
try:
owner_id = _get_owner_id(tab_state_id)
if owner_id is None:
return Response(status=404)
@ -126,6 +132,9 @@ class TabStateView(BaseSupersetView):
).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,6 +155,7 @@ class TabStateView(BaseSupersetView):
@has_access_api
@expose("<int:tab_state_id>/activate", methods=("POST",))
def activate(self, tab_state_id: int) -> FlaskResponse:
try:
owner_id = _get_owner_id(tab_state_id)
if owner_id is None:
return Response(status=404)
@ -159,6 +169,9 @@ class TabStateView(BaseSupersetView):
)
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,18 +182,19 @@ class TabStateView(BaseSupersetView):
if owner_id != get_user_id():
return Response(status=403)
try:
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))
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:
try:
owner_id = _get_owner_id(tab_state_id)
if owner_id is None:
return Response(status=404)
@ -193,10 +207,14 @@ class TabStateView(BaseSupersetView):
)
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:
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
@ -227,12 +245,16 @@ class TabStateView(BaseSupersetView):
).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:
try:
table = json.loads(request.form["table"])
# delete any existing table schema
@ -256,15 +278,22 @@ class TableSchemaView(BaseSupersetView):
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
)
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",))

View File

@ -1013,7 +1013,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)
assert response.status_code == 400
assert response.json["error"] == "Bad request"
# generate query
db.session.add(Query(client_id=client_id, database_id=1))
db.session.commit()