fix(impala): end SQL Lab runs that hit a time limit instead of hanging

docs/admin_docs/configuration/configuring-superset.mdx

@@ -833,7 +833,21 @@ for results, and propagates asynchronous database errors. If an operation is sti
 pending when result fetching begins, the additional polling wait is bounded by
 `SQLLAB_TIMEOUT` (in seconds); on expiry, Superset requests cancellation and reports
 a timeout. Configure the driver's network timeout separately to bound individual
-RPC calls.
+RPC calls. SQL Lab time limits that fire while polling Impala status or logs, or
+while fetching rows, also request cancellation before the timeout propagates,
+including after execution has finished. A time limit can interrupt an RPC
+partway through its response, so before cancelling, Superset closes the impyla
+Thrift connection, empties the read and write buffers of its transport layers
+(buffered, SASL and HTTP), and opens a new connection; the cancellation RPC then
+reads only its own response. Superset then releases the operation and cursor
+handles, also when Impala rejects the cancellation. If the transport cannot be
+reset, or a cleanup RPC fails before reading its whole response (for example on
+a dropped connection or a second time limit), a warning is logged and no
+further cancellation or cursor cleanup RPC is sent; when its transport layers
+are recognised, the connection is also closed with its buffers emptied. The
+operation is then left to Impala's own query and session timeouts. These failures do not hide the timeout. For Hive, a time
+limit while polling or fetching closes the PyHive connection without sending
+cleanup RPCs on it.
 
 ## SQL Lab Query History Pruning
 

superset/db_engine_specs/base.py

@@ -41,6 +41,7 @@
 import pandas as pd
 from apispec import APISpec
 from apispec.ext.marshmallow import MarshmallowPlugin
+from celery.exceptions import SoftTimeLimitExceeded
 from deprecation import deprecated
 from flask import current_app as app, g
 from flask_appbuilder.security.sqla.models import User
@@ -77,6 +78,7 @@
     OAuth2TokenRefreshError,
     SupersetGenericDBErrorException,
     SupersetParseError,
+    SupersetTimeoutException,
 )
 from superset.key_value.types import JsonKeyValueCodec, KeyValueResource
 from superset.sql.parse import (
@@ -134,6 +136,13 @@
 # that we know this is a necessary evil we create an alias, and catch it instead.
 GenericDBException = Exception
 
+# SQL Lab's Celery soft limit for asynchronous runs and ``SQLLAB_TIMEOUT``
+# SIGALRM for main-thread synchronous runs reach the cancellation handlers.
+# The alarm is main-thread only: threaded synchronous requests are not covered
+# by this change. Each limit fires once, so swallowing it leaves the query running
+# with nothing left to stop it.
+TIME_LIMIT_EXCEPTIONS = (SoftTimeLimitExceeded, SupersetTimeoutException)
+
 
 def convert_inspector_columns(cols: list[SQLAColumnType]) -> list[ResultSetColumnType]:
     result_set_columns: list[ResultSetColumnType] = []

superset/db_engine_specs/hive.py

@@ -39,7 +39,11 @@
 from superset import db
 from superset.common.db_query_status import QueryStatus
 from superset.constants import TimeGrain
-from superset.db_engine_specs.base import BaseEngineSpec, DatabaseCategory
+from superset.db_engine_specs.base import (
+    BaseEngineSpec,
+    DatabaseCategory,
+    TIME_LIMIT_EXCEPTIONS,
+)
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.exceptions import SupersetException
 from superset.extensions import cache_manager
@@ -198,15 +202,39 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ..
         import pyhive
         from TCLIService import ttypes
 
-        state = cursor.poll()
-        if state.operationState == ttypes.TOperationState.ERROR_STATE:
-            raise Exception(  # pylint: disable=broad-exception-raised
-                "Query error", state.errorMessage
-            )
         try:
-            return super().fetch_data(cursor, limit)
-        except pyhive.exc.ProgrammingError:
-            return []
+            state = cursor.poll()
+            if state.operationState == ttypes.TOperationState.ERROR_STATE:
+                raise Exception(  # pylint: disable=broad-exception-raised
+                    "Query error", state.errorMessage
+                )
+            try:
+                return super().fetch_data(cursor, limit)
+            except pyhive.exc.ProgrammingError:
+                return []
+        except TIME_LIMIT_EXCEPTIONS:
+            cls._abandon_cursor(cursor)
+            raise
+
+    @staticmethod
+    def _abandon_cursor(cursor: Any) -> None:
+        """
+        Send no further RPC on a transport that may hold an unread reply.
+
+        A time limit raised by a signal can land partway through a reply, and
+        pyhive would read the next reply from the unread rest of it. The
+        cursor's ``close()`` sends CloseOperation only for a live handle.
+        PyHive's connection close unconditionally sends CloseSession, so replace
+        it with a local transport close before attempting to drop the socket.
+        This also prevents RPCs if dropping the socket is itself interrupted.
+        """
+        cursor._operationHandle = None  # pylint: disable=protected-access
+        connection = cursor._connection  # pylint: disable=protected-access
+        connection.close = connection._transport.close  # pylint: disable=protected-access
+        try:
+            connection.close()
+        except Exception:  # pylint: disable=broad-except
+            logger.warning("Closing the transport after a time limit failed")
 
     @classmethod
     def df_to_sql(
@@ -429,10 +457,19 @@ def get_tracking_url_from_logs(cls, log_lines: list[str]) -> str | None:
         return None
 
     @classmethod
-    def handle_cursor(  # pylint: disable=too-many-locals  # noqa: C901
+    def handle_cursor(cls, cursor: Any, query: Query) -> None:
+        """Updates progress information"""
+        try:
+            cls._poll_cursor(cursor, query)
+        except TIME_LIMIT_EXCEPTIONS:
+            cls._abandon_cursor(cursor)
+            raise
+
+    @classmethod
+    def _poll_cursor(  # pylint: disable=too-many-locals  # noqa: C901
         cls, cursor: Any, query: Query
     ) -> None:
-        """Updates progress information"""
+        """Poll the operation until it finishes, recording its progress."""
         # pylint: disable=import-outside-toplevel
         from pyhive import hive
 
@@ -458,6 +495,8 @@ def handle_cursor(  # pylint: disable=too-many-locals  # noqa: C901
             try:
                 logs = cursor.fetch_logs()
                 log = "\n".join(logs) if logs else ""
+            except TIME_LIMIT_EXCEPTIONS:
+                raise
             except Exception:  # pylint: disable=broad-except
                 logger.warning("Call to GetLog() failed")
                 log = ""