Stop button for queries doesn't work in SQL Lab

[DianaLyM]

Bug description

Connected to Apache Kyuubi, stop request returned: {“message”:“Could not cancel query”}

How to reproduce the bug

When I click stop in SQL Lab, nothing happens.

Screenshots/recordings

Superset version

4.0.1

Python version

3.9

Node version

16

Browser

Chrome

Additional context

No response

Checklist

• I have searched Superset docs and Slack and didn't find a solution to my problem.
• I have searched the GitHub issue tracker and didn't find a similar bug report.
• I have checked Superset's logs for errors and if I found a relevant Python stacktrace, I included it here as text in the "additional context" section.

[dosubot]

Hi @DianaLyM! I'm here to help you with bugs, questions, and becoming a contributor. Let's work on this together while we wait for a human maintainer.

To resolve the issue where the stop button for queries in SQL Lab does not work and results in the error message {"message":"Could not cancel query"} when connected to Apache Kyuubi, you need to ensure that the cancel_query function is properly implemented for the Kyuubi engine in the db_engine_spec.

The cancel_query function attempts to cancel a running query by using a cancel query ID stored in the query's extra JSON. If the cancel query ID is not set or the database engine does not support explicit query cancellation, the function will return False.

Here is a relevant snippet from the cancel_query function:

def cancel_query(query: Query) -> bool:
    """
    Cancel a running query.

    Note some engines implicitly handle the cancelation of a query and thus no explicit
    action is required.

    :param query: Query to cancel
    :return: True if query cancelled successfully, False otherwise
    """

    if query.database.db_engine_spec.has_implicit_cancel():
        return True

    # Some databases may need to make preparations for query cancellation
    query.database.db_engine_spec.prepare_cancel_query(query)

    if query.extra.get(QUERY_EARLY_CANCEL_KEY):
        # Query has been cancelled prior to being able to set the cancel key.
        # This can happen if the query cancellation key can only be acquired after the
        # query has been executed
        return True

    cancel_query_id = query.extra.get(QUERY_CANCEL_KEY)
    if cancel_query_id is None:
        return False

    with query.database.get_sqla_engine(
        catalog=query.catalog,
        schema=query.schema,
        source=QuerySource.SQL_LAB,
    ) as engine:
        with closing(engine.raw_connection()) as conn:
            with closing(conn.cursor()) as cursor:
                return query.database.db_engine_spec.cancel_query(
                    cursor, query, cancel_query_id
                )

Additionally, ensure that the stop_query method in the QueryDAO class is correctly implemented. The stop_query method should handle the query status and call the cancel_query function appropriately. Here is the relevant part of the stop_query method:

@staticmethod
def stop_query(client_id: str) -> None:
    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,
        QueryStatus.TIMED_OUT,
    ]:
        logger.warning(
            "Query with client_id could not be stopped: query already complete",
        )
        return

    if not sql_lab.cancel_query(query):
        raise SupersetCancelQueryException("Could not cancel query")

    query.status = QueryStatus.STOPPED
    query.end_time = now_as_float()

To fix the issue, ensure that the db_engine_spec for Apache Kyuubi correctly implements the cancel_query method and handles the cancel_query_id appropriately. If the Kyuubi engine does not support explicit query cancellation, you may need to implement a workaround or ensure that the cancel_query_id is correctly set and used.

Additionally, in the frontend code, ensure that the postStopQuery function in superset-frontend/src/SqlLab/actions/sqlLab.js is correctly implemented to handle the stopping of queries. Here is the relevant part of the postStopQuery function:

export function postStopQuery(query) {
  return function (dispatch) {
    return SupersetClient.post({
      endpoint: '/api/v1/query/stop',
      body: JSON.stringify({ client_id: query.id }),
      headers: { 'Content-Type': 'application/json' },
    })
      .then(() => dispatch(stopQuery(query)))
      .then(() => dispatch(addSuccessToast(t('Query was stopped.'))))
      .catch(() =>
        dispatch(addDangerToast(t('Failed at stopping query. %s', query.id))),
      );
  };
}

Ensure that the stopQuery action is dispatched correctly and that the backend endpoint /api/v1/query/stop is properly handling the stop request ^[1][2][3].

To continue talking to Dosu, mention @dosu.