Skip to content

Commit

Permalink
fix: Exception handling for SQL Lab views (#30897)
Browse files Browse the repository at this point in the history
  • Loading branch information
@michael-s-molina
michael-s-molina authored Nov 12, 2024
1 parent 0e165c1 commit c2885a1
Show file tree
Hide file tree
Showing 2 changed files with 157 additions and 129 deletions.
285 changes: 157 additions & 128 deletions superset/views/sql_lab/views.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,12 @@
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,
)
Expand Down Expand Up @@ -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",))
Expand All @@ -146,19 +155,23 @@ def get(self, tab_state_id: int) -> FlaskResponse:
@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",))
Expand All @@ -169,102 +182,118 @@ def put(self, tab_state_id: int) -> FlaskResponse:
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",))
Expand Down
1 change: 0 additions & 1 deletion tests/integration_tests/core_tests.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1013,7 +1013,6 @@ def test_tabstate_update(self):
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()
Expand Down

0 comments on commit c2885a1

Please sign in to comment.