From f6c8e89b62d1662035ea389e487f0e3efdf11ca8 Mon Sep 17 00:00:00 2001 From: Diego Medina Date: Mon, 9 Jan 2023 16:45:35 -0300 Subject: [PATCH] PR comments --- superset/exceptions.py | 4 ++ superset/queries/api.py | 10 +++-- superset/queries/dao.py | 7 +++- tests/integration_tests/queries/api_tests.py | 29 ++++++++++++- tests/unit_tests/dao/queries_test.py | 44 +++++++++++++++++++- 5 files changed, 87 insertions(+), 7 deletions(-) diff --git a/superset/exceptions.py b/superset/exceptions.py index 153d7439eb790..963bf966820d5 100644 --- a/superset/exceptions.py +++ b/superset/exceptions.py @@ -266,3 +266,7 @@ def __init__(self, error: ValidationError): class SupersetCancelQueryException(SupersetException): status = 422 + + +class QueryNotFoundException(SupersetException): + status = 404 diff --git a/superset/queries/api.py b/superset/queries/api.py index 2639c24becef8..51ba148603285 100644 --- a/superset/queries/api.py +++ b/superset/queries/api.py @@ -23,6 +23,7 @@ from superset import db, event_logger from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod from superset.databases.filters import DatabaseFilter +from superset.exceptions import SupersetException from superset.models.sql_lab import Query from superset.queries.dao import QueryDAO from superset.queries.filters import QueryFilter @@ -190,6 +191,9 @@ def stop_query(self) -> FlaskResponse: 500: $ref: '#/components/responses/500' """ - body = self.stop_query_schema.load(request.json) - QueryDAO.stop_query(body["client_id"]) - return self.response(200, result="OK") + try: + body = self.stop_query_schema.load(request.json) + QueryDAO.stop_query(body["client_id"]) + return self.response(200, result="OK") + except SupersetException as ex: + return self.response(ex.status, message=ex.message) diff --git a/superset/queries/dao.py b/superset/queries/dao.py index e22ed2bcee27e..5867b2917dba0 100644 --- a/superset/queries/dao.py +++ b/superset/queries/dao.py @@ -21,7 +21,7 @@ from superset import sql_lab from superset.common.db_query_status import QueryStatus from superset.dao.base import BaseDAO -from superset.exceptions import SupersetCancelQueryException +from superset.exceptions import QueryNotFoundException, SupersetCancelQueryException from superset.extensions import db from superset.models.sql_lab import Query, SavedQuery from superset.queries.filters import QueryFilter @@ -63,7 +63,10 @@ def save_metadata(query: Query, payload: Dict[str, Any]) -> None: @staticmethod def stop_query(client_id: str) -> None: - query = db.session.query(Query).filter_by(client_id=client_id).one() + query = db.session.query(Query).filter_by(client_id=client_id).one_or_none() + if not query: + raise QueryNotFoundException(f"Query with client_id {client_id} not found") + if query.status in [ QueryStatus.FAILED, QueryStatus.SUCCESS, diff --git a/tests/integration_tests/queries/api_tests.py b/tests/integration_tests/queries/api_tests.py index fe6d0d909072d..8c193662a12c7 100644 --- a/tests/integration_tests/queries/api_tests.py +++ b/tests/integration_tests/queries/api_tests.py @@ -394,6 +394,31 @@ def test_get_list_query_no_data_access(self): db.session.delete(query) db.session.commit() + @mock.patch("superset.sql_lab.cancel_query") + @mock.patch("superset.views.core.db.session") + def test_stop_query_not_found( + self, mock_superset_db_session, mock_sql_lab_cancel_query + ): + """ + Handles stop query when the DB engine spec does not + have a cancel query method (with invalid client_id). + """ + form_data = {"client_id": "foo2"} + query_mock = mock.Mock() + query_mock.return_value = None + self.login(username="admin") + mock_superset_db_session.query().filter_by().one_or_none = query_mock + mock_sql_lab_cancel_query.return_value = True + rv = self.client.post( + "/api/v1/query/stop", + data=json.dumps(form_data), + content_type="application/json", + ) + + assert rv.status_code == 404 + data = json.loads(rv.data.decode("utf-8")) + assert data["message"] == "Query with client_id foo2 not found" + @mock.patch("superset.sql_lab.cancel_query") @mock.patch("superset.views.core.db.session") def test_stop_query(self, mock_superset_db_session, mock_sql_lab_cancel_query): @@ -406,7 +431,9 @@ def test_stop_query(self, mock_superset_db_session, mock_sql_lab_cancel_query): query_mock.client_id = "foo" query_mock.status = QueryStatus.RUNNING self.login(username="admin") - mock_superset_db_session.query().filter_by().one().return_value = query_mock + mock_superset_db_session.query().filter_by().one_or_none().return_value = ( + query_mock + ) mock_sql_lab_cancel_query.return_value = True rv = self.client.post( "/api/v1/query/stop", diff --git a/tests/unit_tests/dao/queries_test.py b/tests/unit_tests/dao/queries_test.py index 715811c60387f..590ba92d48f13 100644 --- a/tests/unit_tests/dao/queries_test.py +++ b/tests/unit_tests/dao/queries_test.py @@ -21,7 +21,7 @@ from pytest_mock import MockFixture from sqlalchemy.orm.session import Session -from superset.exceptions import SupersetCancelQueryException +from superset.exceptions import QueryNotFoundException, SupersetCancelQueryException def test_query_dao_save_metadata(session: Session) -> None: @@ -58,6 +58,48 @@ def test_query_dao_save_metadata(session: Session) -> None: assert query.extra.get("columns", None) == [] +def test_query_dao_stop_query_not_found( + mocker: MockFixture, app: Any, session: Session +) -> None: + from superset.common.db_query_status import QueryStatus + from superset.models.core import Database + from superset.models.sql_lab import Query + + engine = session.get_bind() + Query.metadata.create_all(engine) # pylint: disable=no-member + + db = Database(database_name="my_database", sqlalchemy_uri="sqlite://") + + query_obj = Query( + client_id="foo", + database=db, + tab_name="test_tab", + sql_editor_id="test_editor_id", + sql="select * from bar", + select_sql="select * from bar", + executed_sql="select * from bar", + limit=100, + select_as_cta=False, + rows=100, + error_message="none", + results_key="abc", + status=QueryStatus.RUNNING, + ) + + session.add(db) + session.add(query_obj) + + mocker.patch("superset.sql_lab.cancel_query", return_value=False) + + from superset.queries.dao import QueryDAO + + with pytest.raises(QueryNotFoundException): + QueryDAO.stop_query("foo2") + + query = session.query(Query).one() + assert query.status == QueryStatus.RUNNING + + def test_query_dao_stop_query_not_running( mocker: MockFixture, app: Any, session: Session ) -> None: