fix(trino): cancel_query doesn't abort query execution in Trino server

Author: dungdm93

superset/db_engine_specs/trino.py

@@ -134,6 +134,19 @@ def get_url_for_impersonation(
     def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool:
         return True
 
+    @classmethod
+    def execute(cls, cursor: Cursor, query: str, **kwargs: Any) -> None:
+        opts = {}
+        if "async_" in kwargs:
+            opts["deferred_fetch"] = kwargs["async_"]
+        if cls.arraysize:
+            cursor.arraysize = cls.arraysize
+
+        try:
+            cursor.execute(query, **opts)
+        except Exception as ex:
+            raise cls.get_dbapi_mapped_exception(ex)
+
     @classmethod
     def get_tracking_url(cls, cursor: Cursor) -> str | None:
         try:
@@ -160,8 +173,8 @@ def handle_cursor(cls, cursor: Cursor, query: Query, session: Session) -> None:
 
         session.commit()
 
-        # if query cancelation was requested prior to the handle_cursor call, but
-        # the query was still executed, trigger the actual query cancelation now
+        # if query cancellation was requested prior to the handle_cursor call, but
+        # the query was still executed, trigger the actual query cancellation now
         if query.extra.get(QUERY_EARLY_CANCEL_KEY):
             cls.cancel_query(
                 cursor=cursor,
@@ -178,7 +191,7 @@ def prepare_cancel_query(cls, query: Query, session: Session) -> None:
             session.commit()
 
     @classmethod
-    def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool:
+    def cancel_query(cls, cursor: Cursor, query: Query, cancel_query_id: str) -> bool:
         """
         Cancel query in the underlying database.
 
@@ -188,11 +201,12 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool:
         :return: True if query cancelled successfully, False otherwise
         """
         try:
+            # Can't use cursor.cancel() because
+            # cursor is new object created in sql_lab.cancel_query
             cursor.execute(
                 f"CALL system.runtime.kill_query(query_id => '{cancel_query_id}',"
                 "message => 'Query cancelled by Superset')"
             )
-            cursor.fetchall()  # needed to trigger the call
         except Exception:  # pylint: disable=broad-except
             return False