PR comments

superset/exceptions.py

@@ -266,3 +266,7 @@ def __init__(self, error: ValidationError):
 
 class SupersetCancelQueryException(SupersetException):
     status = 422
+
+
+class QueryNotFoundException(SupersetException):
+    status = 404

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)

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,

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",

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: