fix(15403): Re-enable canceling query for Hive and Presto (#15878)

Co-authored-by: John Bodley <john.bodley@airbnb.com>

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

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):
 

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

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
+            )

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)