diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 143eb8680e..1ce082c4f6 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1305,6 +1305,18 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods ) return None + @classmethod + def has_implicit_cancel(cls) -> bool: + """ + Return True if the live cursor handles the implicit cancelation of the query, + False otherise. + + :return: Whether the live cursor implicitly cancels the query + :see: handle_cursor + """ + + return False + @classmethod def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: """ @@ -1316,6 +1328,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods :param query: Query instance :return: Query identifier """ + return None @classmethod @@ -1330,6 +1343,8 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods :return: True if query cancelled successfully, False otherwise """ + return False + # schema for adding a database by providing parameters instead of the # full SQLAlchemy URI diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 73cfb0e711..39152528e8 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -546,6 +546,18 @@ class HiveEngineSpec(PrestoEngineSpec): or parsed_query.is_show() ) + @classmethod + def has_implicit_cancel(cls) -> bool: + """ + Return True if the live cursor handles the implicit cancelation of the query, + False otherise. + + :return: Whether the live cursor implicitly cancels the query + :see: handle_cursor + """ + + return True + class SparkEngineSpec(HiveEngineSpec): diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 27bc98fef4..d3567e98c1 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -1234,3 +1234,15 @@ class PrestoEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-metho return column_spec return super().get_column_spec(native_type) + + @classmethod + def has_implicit_cancel(cls) -> bool: + """ + Return True if the live cursor handles the implicit cancelation of the query, + False otherise. + + :return: Whether the live cursor implicitly cancels the query + :see: handle_cursor + """ + + return True diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 1ebf10e599..d58124d73d 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -587,23 +587,30 @@ def cancel_query(query: Query, user_name: Optional[str] = None) -> bool: """ Cancel a running query. + Note some engines implicitly handle the cancelation of a query and thus no expliicit + action is required. + :param query: Query to cancel :param user_name: Default username :return: True if query cancelled successfully, False otherwise """ - cancel_query_id = query.extra.get(cancel_query_key, None) + + if query.database.db_engine_spec.has_implicit_cancel(): + return True + + cancel_query_id = query.extra.get(cancel_query_key) if cancel_query_id is None: return False - database = query.database - engine = database.get_sqla_engine( + engine = query.database.get_sqla_engine( schema=query.schema, nullpool=True, user_name=user_name, source=QuerySource.SQL_LAB, ) - db_engine_spec = database.db_engine_spec with closing(engine.raw_connection()) as conn: with closing(conn.cursor()) as cursor: - return db_engine_spec.cancel_query(cursor, query, cancel_query_id) + return query.database.db_engine_spec.cancel_query( + cursor, query, cancel_query_id + ) diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py index 442001fe10..f5b7859ba9 100644 --- a/tests/integration_tests/sqllab_tests.py +++ b/tests/integration_tests/sqllab_tests.py @@ -30,12 +30,15 @@ import prison from superset import db, security_manager from superset.connectors.sqla.models import SqlaTable from superset.db_engine_specs import BaseEngineSpec +from superset.db_engine_specs.hive import HiveEngineSpec +from superset.db_engine_specs.presto import PrestoEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import SupersetErrorException from superset.models.core import Database from superset.models.sql_lab import LimitingFactor, Query, SavedQuery from superset.result_set import SupersetResultSet from superset.sql_lab import ( + cancel_query, execute_sql_statements, execute_sql_statement, get_sql_results, @@ -985,3 +988,10 @@ class TestSqlLab(SupersetTestCase): } ] } + + +@pytest.mark.parametrize("spec", [HiveEngineSpec, PrestoEngineSpec]) +def test_cancel_query_implicit(spec: BaseEngineSpec) -> None: + query = mock.MagicMock() + query.database.db_engine_spec = spec + assert cancel_query(query)