iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

275
superset/views/sql_lab/views.py
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.models.sql_lab import Query, SavedQuery, TableSchema, TabState
from superset.superset_typing import FlaskResponse from superset.superset_typing import FlaskResponse
from superset.utils import json 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 ( from superset.views.base import (
BaseSupersetView, BaseSupersetView,
DeleteMixin, DeleteMixin,
DeprecateModelViewMixin, DeprecateModelViewMixin,
json_error_response,
json_success, json_success,
SupersetModelView, SupersetModelView,
) )
@ -84,48 +85,56 @@ class TabStateView(BaseSupersetView):
@has_access_api @has_access_api
@expose("/", methods=("POST",)) @expose("/", methods=("POST",))
def post(self) -> FlaskResponse: def post(self) -> FlaskResponse:
query_editor = json.loads(request.form["queryEditor"]) try:
tab_state = TabState( query_editor = json.loads(request.form["queryEditor"])
user_id=get_user_id(), tab_state = TabState(
# This is for backward compatibility user_id=get_user_id(),
label=query_editor.get("name") # This is for backward compatibility
or query_editor.get("title", __("Untitled Query")), label=query_editor.get("name")
active=True, or query_editor.get("title", __("Untitled Query")),
database_id=query_editor["dbId"], active=True,
catalog=query_editor.get("catalog"), database_id=query_editor["dbId"],
schema=query_editor.get("schema"), catalog=query_editor.get("catalog"),
sql=query_editor.get("sql", "SELECT ..."), schema=query_editor.get("schema"),
query_limit=query_editor.get("queryLimit"), sql=query_editor.get("sql", "SELECT ..."),
hide_left_bar=query_editor.get("hideLeftBar"), query_limit=query_editor.get("queryLimit"),
saved_query_id=query_editor.get("remoteId"), hide_left_bar=query_editor.get("hideLeftBar"),
template_params=query_editor.get("templateParams"), saved_query_id=query_editor.get("remoteId"),
) template_params=query_editor.get("templateParams"),
( )
db.session.query(TabState) (
.filter_by(user_id=get_user_id()) db.session.query(TabState)
.update({"active": False}) .filter_by(user_id=get_user_id())
) .update({"active": False})
db.session.add(tab_state) )
db.session.commit() db.session.add(tab_state)
return json_success(json.dumps({"id": tab_state.id})) 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 @has_access_api
@expose("/<int:tab_state_id>", methods=("DELETE",)) @expose("/<int:tab_state_id>", methods=("DELETE",))
def delete(self, tab_state_id: int) -> FlaskResponse: def delete(self, tab_state_id: int) -> FlaskResponse:
owner_id = _get_owner_id(tab_state_id) try:
if owner_id is None: owner_id = _get_owner_id(tab_state_id)
return Response(status=404) if owner_id is None:
if owner_id != get_user_id(): return Response(status=404)
return Response(status=403) if owner_id != get_user_id():
return Response(status=403)
db.session.query(TabState).filter(TabState.id == tab_state_id).delete( db.session.query(TabState).filter(TabState.id == tab_state_id).delete(
synchronize_session=False synchronize_session=False
) )
db.session.query(TableSchema).filter( db.session.query(TableSchema).filter(
TableSchema.tab_state_id == tab_state_id TableSchema.tab_state_id == tab_state_id
).delete(synchronize_session=False) ).delete(synchronize_session=False)
db.session.commit() db.session.commit()
return json_success(json.dumps("OK")) 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 @has_access_api
@expose("/<int:tab_state_id>", methods=("GET",)) @expose("/<int:tab_state_id>", methods=("GET",))
@ -146,19 +155,23 @@ class TabStateView(BaseSupersetView):
@has_access_api @has_access_api
@expose("<int:tab_state_id>/activate", methods=("POST",)) @expose("<int:tab_state_id>/activate", methods=("POST",))
def activate(self, tab_state_id: int) -> FlaskResponse: def activate(self, tab_state_id: int) -> FlaskResponse:
owner_id = _get_owner_id(tab_state_id) try:
if owner_id is None: owner_id = _get_owner_id(tab_state_id)
return Response(status=404) if owner_id is None:
if owner_id != get_user_id(): return Response(status=404)
return Response(status=403) if owner_id != get_user_id():
return Response(status=403)
( (
db.session.query(TabState) db.session.query(TabState)
.filter_by(user_id=get_user_id()) .filter_by(user_id=get_user_id())
.update({"active": TabState.id == tab_state_id}) .update({"active": TabState.id == tab_state_id})
) )
db.session.commit() db.session.commit()
return json_success(json.dumps(tab_state_id)) 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 @has_access_api
@expose("<int:tab_state_id>", methods=("PUT",)) @expose("<int:tab_state_id>", methods=("PUT",))
@ -169,102 +182,118 @@ class TabStateView(BaseSupersetView):
if owner_id != get_user_id(): if owner_id != get_user_id():
return Response(status=403) return Response(status=403)
fields = {k: json.loads(v) for k, v in request.form.to_dict().items()} try:
if client_id := fields.get("latest_query_id"): fields = {k: json.loads(v) for k, v in request.form.to_dict().items()}
query = db.session.query(Query).filter_by(client_id=client_id).one_or_none() db.session.query(TabState).filter_by(id=tab_state_id).update(fields)
if not query: db.session.commit()
return self.json_response({"error": "Bad request"}, status=400) return json_success(json.dumps(tab_state_id))
db.session.query(TabState).filter_by(id=tab_state_id).update(fields) except Exception as ex: # pylint: disable=broad-except
db.session.commit() db.session.rollback()
return json_success(json.dumps(tab_state_id)) return json_error_response(error_msg_from_exception(ex), 400)
@has_access_api @has_access_api
@expose("<int:tab_state_id>/migrate_query", methods=("POST",)) @expose("<int:tab_state_id>/migrate_query", methods=("POST",))
def migrate_query(self, tab_state_id: int) -> FlaskResponse: def migrate_query(self, tab_state_id: int) -> FlaskResponse:
owner_id = _get_owner_id(tab_state_id) try:
if owner_id is None: owner_id = _get_owner_id(tab_state_id)
return Response(status=404) if owner_id is None:
if owner_id != get_user_id(): return Response(status=404)
return Response(status=403) if owner_id != get_user_id():
return Response(status=403)
client_id = json.loads(request.form["queryId"]) client_id = json.loads(request.form["queryId"])
db.session.query(Query).filter_by(client_id=client_id).update( db.session.query(Query).filter_by(client_id=client_id).update(
{"sql_editor_id": tab_state_id} {"sql_editor_id": tab_state_id}
) )
db.session.commit() db.session.commit()
return json_success(json.dumps(tab_state_id)) 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 @has_access_api
@expose("<int:tab_state_id>/query/<client_id>", methods=("DELETE",)) @expose("<int:tab_state_id>/query/<client_id>", methods=("DELETE",))
def delete_query(self, tab_state_id: int, client_id: str) -> FlaskResponse: def delete_query(self, tab_state_id: int, client_id: str) -> FlaskResponse:
# Before deleting the query, ensure it's not tied to any try:
# active tab as the last query. If so, replace the query # Before deleting the query, ensure it's not tied to any
# with the latest one created in that tab # active tab as the last query. If so, replace the query
tab_state_query = db.session.query(TabState).filter_by( # with the latest one created in that tab
id=tab_state_id, latest_query_id=client_id tab_state_query = db.session.query(TabState).filter_by(
) id=tab_state_id, latest_query_id=client_id
if tab_state_query.count(): )
query = ( if tab_state_query.count():
db.session.query(Query) query = (
.filter( db.session.query(Query)
and_( .filter(
Query.client_id != client_id, and_(
Query.user_id == get_user_id(), Query.client_id != client_id,
Query.sql_editor_id == str(tab_state_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}
) )
.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( db.session.query(Query).filter_by(
client_id=client_id, client_id=client_id,
user_id=get_user_id(), user_id=get_user_id(),
sql_editor_id=str(tab_state_id), sql_editor_id=str(tab_state_id),
).delete(synchronize_session=False) ).delete(synchronize_session=False)
db.session.commit() db.session.commit()
return json_success(json.dumps("OK")) 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): class TableSchemaView(BaseSupersetView):
@has_access_api @has_access_api
@expose("/", methods=("POST",)) @expose("/", methods=("POST",))
def post(self) -> FlaskResponse: def post(self) -> FlaskResponse:
table = json.loads(request.form["table"]) try:
table = json.loads(request.form["table"])
# delete any existing table schema # delete any existing table schema
db.session.query(TableSchema).filter( db.session.query(TableSchema).filter(
TableSchema.tab_state_id == table["queryEditorId"], TableSchema.tab_state_id == table["queryEditorId"],
TableSchema.database_id == table["dbId"], TableSchema.database_id == table["dbId"],
TableSchema.catalog == table.get("catalog"), TableSchema.catalog == table.get("catalog"),
TableSchema.schema == table["schema"], TableSchema.schema == table["schema"],
TableSchema.table == table["name"], TableSchema.table == table["name"],
).delete(synchronize_session=False) ).delete(synchronize_session=False)
table_schema = TableSchema( table_schema = TableSchema(
tab_state_id=table["queryEditorId"], tab_state_id=table["queryEditorId"],
database_id=table["dbId"], database_id=table["dbId"],
catalog=table.get("catalog"), catalog=table.get("catalog"),
schema=table["schema"], schema=table["schema"],
table=table["name"], table=table["name"],
description=json.dumps(table), description=json.dumps(table),
expanded=True, expanded=True,
) )
db.session.add(table_schema) db.session.add(table_schema)
db.session.commit() db.session.commit()
return json_success(json.dumps({"id": table_schema.id})) 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 @has_access_api
@expose("/<int:table_schema_id>", methods=("DELETE",)) @expose("/<int:table_schema_id>", methods=("DELETE",))
def delete(self, table_schema_id: int) -> FlaskResponse: def delete(self, table_schema_id: int) -> FlaskResponse:
db.session.query(TableSchema).filter(TableSchema.id == table_schema_id).delete( try:
synchronize_session=False db.session.query(TableSchema).filter(
) TableSchema.id == table_schema_id
db.session.commit() ).delete(synchronize_session=False)
return json_success(json.dumps("OK")) 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 @has_access_api
@expose("/<int:table_schema_id>/expanded", methods=("POST",)) @expose("/<int:table_schema_id>/expanded", methods=("POST",))

1
tests/integration_tests/core_tests.py
View File

@ -1013,7 +1013,6 @@ class TestCore(SupersetTestCase):
data = {"sql": json.dumps("select 1"), "latest_query_id": json.dumps(client_id)} data = {"sql": json.dumps("select 1"), "latest_query_id": json.dumps(client_id)}
response = self.client.put(f"/tabstateview/{tab_state_id}", data=data) response = self.client.put(f"/tabstateview/{tab_state_id}", data=data)
assert response.status_code == 400 assert response.status_code == 400
assert response.json["error"] == "Bad request"
# generate query # generate query
db.session.add(Query(client_id=client_id, database_id=1)) db.session.add(Query(client_id=client_id, database_id=1))
db.session.commit() db.session.commit()