From ec620335d67ca988ab49b7078b823533298545a0 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 29 Oct 2026 13:16:39 +0000 Subject: [PATCH 01/14] fix(impala): end SQL Lab runs that hit a time limit instead of hanging A SQL Lab time limit is raised once into whatever the worker is executing: the synchronous SQLLAB_TIMEOUT alarm raises SupersetTimeoutException and the Celery soft limit raises SoftTimeLimitExceeded. Impala's handle_cursor only treated the soft limit specially, so: - a synchronous timeout raised while polling fell into the broad `except Exception` and was swallowed, and the Run kept waiting for the whole query before failing from fetch_data's own deadline; - either limit raised while the progress log was fetched was swallowed as a failed GetLog() call, and with the limit already spent nothing stopped the query until the hard limit; - a limit raised while fetch_data waited on the operation left it running on the coordinator. Both limits are now cancelled and propagated from handle_cursor, including from the progress log call and the cancel helper, and fetch_data cancels the operation before propagating them. Hive's handle_cursor had the same swallow around fetch_logs and now lets the limits through. Presto and Trino do not wrap their polling in a broad except. --- superset/db_engine_specs/base.py | 8 ++ superset/db_engine_specs/hive.py | 8 +- superset/db_engine_specs/impala.py | 31 +++-- tests/unit_tests/db_engine_specs/test_hive.py | 51 ++++++++ .../unit_tests/db_engine_specs/test_impala.py | 29 +++++ .../db_engine_specs/test_impala_pending.py | 120 +++++++++++++++++- 6 files changed, 237 insertions(+), 10 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 7eadb4a94d56..a6e4c8a9eb18 100644 --- a/superset/db_engine_specs/base.py +++ b/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,12 @@ # that we know this is a necessary evil we create an alias, and catch it instead. GenericDBException = Exception +# SQL Lab's time limits are raised into whatever the worker is executing: the +# Celery soft limit for asynchronous runs and the ``SQLLAB_TIMEOUT`` alarm for +# synchronous ones. Each fires once, so a handler that swallows 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] = [] diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index b85868ea7735..c1507f4f322b 100644 --- a/superset/db_engine_specs/hive.py +++ b/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 @@ -458,6 +462,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 = "" diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index e5134e21b947..7102402347e5 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -23,14 +23,17 @@ from datetime import datetime from typing import Any, Optional, TYPE_CHECKING -from celery.exceptions import SoftTimeLimitExceeded from flask import current_app as app from sqlalchemy import text, types from sqlalchemy.engine.reflection import Inspector from superset import db from superset.constants import QUERY_EARLY_CANCEL_KEY, TimeGrain -from superset.db_engine_specs.base import BaseEngineSpec, DatabaseCategory +from superset.db_engine_specs.base import ( + BaseEngineSpec, + DatabaseCategory, + TIME_LIMIT_EXCEPTIONS, +) from superset.models.sql_lab import Query from superset.utils.core import QueryStatus from superset.utils.network import get_ssrf_safe_requester, is_safe_host @@ -152,6 +155,16 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. cursor.fetchall() except Error as ex: raise cls.get_dbapi_mapped_exception(ex) from ex + except TIME_LIMIT_EXCEPTIONS: + # The operation is released before the time limit propagates so + # it does not keep running on the coordinator. + try: + cursor.cancel_operation() + except TIME_LIMIT_EXCEPTIONS: + raise + except Exception: # pylint: disable=broad-except + logger.warning("cancel_operation() failed after a time limit") + raise return super().fetch_data(cursor, limit) @classmethod @@ -159,7 +172,7 @@ def _cancel_operation(cls, cursor: Any, query_id: int) -> None: """Cancel the live operation and release its handles.""" try: cursor.cancel_operation() - except SoftTimeLimitExceeded: + except TIME_LIMIT_EXCEPTIONS: raise except Exception: # pylint: disable=broad-except logger.warning("Query %s: cancel_operation() failed", query_id) @@ -168,13 +181,13 @@ def _cancel_operation(cls, cursor: Any, query_id: int) -> None: # of the connection's life. try: cursor.close_operation() - except SoftTimeLimitExceeded: + except TIME_LIMIT_EXCEPTIONS: raise except Exception: # pylint: disable=broad-except logger.warning("Query %s: close_operation() failed", query_id) try: cursor.close() - except SoftTimeLimitExceeded: + except TIME_LIMIT_EXCEPTIONS: raise except Exception: # pylint: disable=broad-except logger.warning("Query %s: close() failed", query_id) @@ -219,6 +232,8 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: if status == "RUNNING_STATE": try: log = cursor.get_log() or "" + except TIME_LIMIT_EXCEPTIONS: + raise except Exception: # pylint: disable=broad-except logger.warning("Call to GetLog() failed") log = "" @@ -234,9 +249,9 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: pending_sleep_interval *= 2 time.sleep(sleep_interval) status = cursor.status() - except SoftTimeLimitExceeded: - # SQL Lab's own handler marks the query TIMED_OUT once this propagates; - # the operation is cancelled first so it does not keep running on the + except TIME_LIMIT_EXCEPTIONS: + # SQL Lab ends the query as timed out once this propagates; the + # operation is cancelled first so it does not keep running on the # coordinator while fetch_data() waits on it. cls._cancel_operation(cursor, query_id) raise diff --git a/tests/unit_tests/db_engine_specs/test_hive.py b/tests/unit_tests/db_engine_specs/test_hive.py index dc6f9b071e1b..a6f61ab3554c 100644 --- a/tests/unit_tests/db_engine_specs/test_hive.py +++ b/tests/unit_tests/db_engine_specs/test_hive.py @@ -18,13 +18,18 @@ from datetime import datetime from typing import Optional +from unittest.mock import Mock, patch import pytest +from celery.exceptions import SoftTimeLimitExceeded +from flask import Flask from pytest_mock import MockerFixture from sqlalchemy.engine.interfaces import Dialect from sqlalchemy.engine.url import make_url from superset.constants import TimeGrain +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import SupersetTimeoutException from superset.sql.parse import Table from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm from tests.unit_tests.fixtures.common import dttm # noqa: F401 @@ -373,3 +378,49 @@ def test_the_partition_transform_default_is_not_inherited_by_presto() -> None: from superset.db_engine_specs.presto import PrestoEngineSpec assert PrestoEngineSpec.partition_value_transform_default is None + + +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +def test_handle_cursor_propagates_time_limits(timeout: Exception) -> None: + """ + A SQL Lab time limit that fires while the logs are fetched is not a failed + ``fetch_logs`` call: it fires only once, so swallowing it leaves the query + running until the hard limit. + """ + from pyhive import hive + + from superset.db_engine_specs.hive import HiveEngineSpec + + app: Flask = Flask(__name__) + app.config["DB_POLL_INTERVAL_SECONDS"] = {"hive": 0} + query: Mock = Mock(id=1, progress=0) + cursor: Mock = Mock() + # Finite, so a swallowed time limit shows up as a failed assertion rather than + # a loop that never ends. + cursor.poll.side_effect = [ + Mock(operationState=hive.ttypes.TOperationState.RUNNING_STATE), + Mock(operationState=hive.ttypes.TOperationState.RUNNING_STATE), + Mock(operationState=hive.ttypes.TOperationState.FINISHED_STATE), + ] + cursor.fetch_logs.side_effect = timeout + with ( + app.app_context(), + patch("superset.db_engine_specs.hive.db") as db, + patch("superset.db_engine_specs.hive.time.sleep"), + ): + db.session.query.return_value.filter_by.return_value.one.return_value = query + query.status = "running" + with pytest.raises(type(timeout)): + HiveEngineSpec.handle_cursor(cursor, query) + cursor.fetch_logs.assert_called_once_with() diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index f8603e2f5410..4020cc7f308a 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -20,9 +20,12 @@ from unittest.mock import Mock, patch import pytest +from celery.exceptions import SoftTimeLimitExceeded from pytest_mock import MockerFixture from superset.db_engine_specs.impala import ImpalaEngineSpec as spec # noqa: N813 +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import SupersetTimeoutException from superset.models.core import Database from superset.models.sql_lab import Query from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm @@ -304,3 +307,29 @@ def test_engines_without_a_sensible_default_offer_none() -> None: assert BaseEngineSpec.partition_value_transform_default is None assert PostgresEngineSpec.partition_value_transform_default is None + + +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_fetch_data_cancels_the_operation_on_time_limit( + sleep_mock: Mock, timeout: Exception, app_context: None, impala_errors: None +) -> None: + """A SQL Lab time limit hitting the wait must not leave the operation running.""" + sleep_mock.side_effect = timeout + cursor = Mock() + cursor.is_executing.return_value = True + with pytest.raises(type(timeout)): + spec.fetch_data(cursor) + cursor.cancel_operation.assert_called_once_with() + cursor.fetchall.assert_not_called() diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 1e7d69b76e78..0ad054a71d3f 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_pending.py +++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py @@ -29,8 +29,15 @@ QUERY_EARLY_CANCEL_KEY, ) from superset.db_engine_specs.impala import ImpalaEngineSpec +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import SupersetTimeoutException from superset.sql.execution.executor import SQLExecutor -from superset.sql_lab import cancel_query, execute_query, SqlLabQueryStoppedException +from superset.sql_lab import ( + cancel_query, + execute_query, + SqlLabException, + SqlLabQueryStoppedException, +) from superset.utils.core import QueryStatus @@ -401,3 +408,114 @@ def refresh_query(refreshed_query: Mock) -> None: cursor.cancel_operation.assert_called_once_with() cursor.close_operation.assert_called_once_with() cursor.close.assert_called_once_with() + + +def _sync_timeout() -> SupersetTimeoutException: + """Build the exception SQL Lab's synchronous ``SQLLAB_TIMEOUT`` alarm raises.""" + return SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ) + + +@pytest.mark.parametrize("state", ["PENDING_STATE", "RUNNING_STATE"]) +def test_sync_timeout_cancels_the_operation(state: str) -> None: + """ + The synchronous ``SQLLAB_TIMEOUT`` firing while the loop polls must cancel the + operation and propagate, rather than being swallowed so the Run keeps waiting + for the whole query. + """ + app: Flask = Flask(__name__) + app.config["DB_POLL_INTERVAL_SECONDS"] = {"impala": 5} + query: Mock = Mock(id=1, extra={}, status="running", progress=0) + cursor: Mock = Mock() + cursor.status.return_value = state + cursor.get_log.return_value = "" + timeout: SupersetTimeoutException = _sync_timeout() + with ( + app.app_context(), + patch("superset.db_engine_specs.impala.db"), + patch("superset.db_engine_specs.impala.time.sleep", side_effect=timeout), + ): + with pytest.raises(SupersetTimeoutException): + ImpalaEngineSpec.handle_cursor(cursor, query) + cursor.cancel_operation.assert_called_once_with() + cursor.close_operation.assert_called_once_with() + cursor.close.assert_called_once_with() + + +@pytest.mark.parametrize( + "timeout", [SoftTimeLimitExceeded(), _sync_timeout()], ids=["async", "sync"] +) +def test_timeout_during_progress_log_cancels_the_operation( + timeout: Exception, +) -> None: + """ + A time limit that fires while the progress log is fetched must not be mistaken + for a failed ``GetLog()`` call: it fires only once, so swallowing it leaves + the query running until the hard limit. + """ + app: Flask = Flask(__name__) + app.config["DB_POLL_INTERVAL_SECONDS"] = {"impala": 5} + query: Mock = Mock(id=1, extra={}, status="running", progress=0) + cursor: Mock = Mock() + # Finite, so a swallowed time limit shows up as a failed assertion rather than + # a loop that never ends. + cursor.status.side_effect = ["RUNNING_STATE", "RUNNING_STATE", "FINISHED_STATE"] + cursor.get_log.side_effect = timeout + with ( + app.app_context(), + patch("superset.db_engine_specs.impala.db"), + patch("superset.db_engine_specs.impala.time.sleep") as sleep, + ): + with pytest.raises(type(timeout)): + ImpalaEngineSpec.handle_cursor(cursor, query) + sleep.assert_not_called() + cursor.cancel_operation.assert_called_once_with() + cursor.close_operation.assert_called_once_with() + cursor.close.assert_called_once_with() + + +def test_status_failure_still_ends_polling_without_raising() -> None: + """Ordinary polling failures keep ending the loop quietly.""" + app: Flask = Flask(__name__) + query: Mock = Mock(id=1, extra={}, status="running", progress=0) + cursor: Mock = Mock() + cursor.status.side_effect = RuntimeError("rpc failed") + with app.app_context(), patch("superset.db_engine_specs.impala.db"): + ImpalaEngineSpec.handle_cursor(cursor, query) + cursor.cancel_operation.assert_not_called() + + +def test_sync_timeout_fails_the_run_and_cancels_the_operation( + mocker: MockerFixture, app_context: None +) -> None: + """A synchronous timeout during polling ends the Run with the timeout message.""" + query = MagicMock( + id=1, + status=QueryStatus.RUNNING, + progress=0, + executed_sql="SELECT 1", + extra={}, + limit=None, + ) + query.database.db_engine_spec = ImpalaEngineSpec + cursor = MagicMock() + cursor.status.return_value = "RUNNING_STATE" + cursor.get_log.return_value = "" + mocker.patch( + "superset.db_engine_specs.impala.time.sleep", side_effect=_sync_timeout() + ) + mocker.patch("superset.db.session") + mocker.patch("superset.sql_lab.warm_and_release_connection") + mocker.patch( + "superset.sql_lab.event_logger.log_context", return_value=nullcontext() + ) + mocker.patch.object(ImpalaEngineSpec, "get_cancel_query_id", return_value=None) + + with pytest.raises(SqlLabException, match="exceeded the 30 seconds timeout"): + execute_query(query, cursor) + + cursor.cancel_operation.assert_called_once_with() + cursor.fetchall.assert_not_called() From 924d043b5b405f0b3e5a5879f83d7b5cbec144ac Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 29 Oct 2026 13:49:50 +0000 Subject: [PATCH 02/14] test(hive): don't require the pyhive thrift dependency in the time-limit test --- tests/unit_tests/db_engine_specs/test_hive.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/db_engine_specs/test_hive.py b/tests/unit_tests/db_engine_specs/test_hive.py index a6f61ab3554c..b000e1facc7d 100644 --- a/tests/unit_tests/db_engine_specs/test_hive.py +++ b/tests/unit_tests/db_engine_specs/test_hive.py @@ -398,10 +398,16 @@ def test_handle_cursor_propagates_time_limits(timeout: Exception) -> None: ``fetch_logs`` call: it fires only once, so swallowing it leaves the query running until the hard limit. """ - from pyhive import hive - from superset.db_engine_specs.hive import HiveEngineSpec + # The pyhive driver needs thrift, which is not always installed. + states = Mock( + INITIALIZED_STATE="INITIALIZED_STATE", + RUNNING_STATE="RUNNING_STATE", + FINISHED_STATE="FINISHED_STATE", + ) + hive = Mock(ttypes=Mock(TOperationState=states)) + app: Flask = Flask(__name__) app.config["DB_POLL_INTERVAL_SECONDS"] = {"hive": 0} query: Mock = Mock(id=1, progress=0) @@ -409,15 +415,16 @@ def test_handle_cursor_propagates_time_limits(timeout: Exception) -> None: # Finite, so a swallowed time limit shows up as a failed assertion rather than # a loop that never ends. cursor.poll.side_effect = [ - Mock(operationState=hive.ttypes.TOperationState.RUNNING_STATE), - Mock(operationState=hive.ttypes.TOperationState.RUNNING_STATE), - Mock(operationState=hive.ttypes.TOperationState.FINISHED_STATE), + Mock(operationState=states.RUNNING_STATE), + Mock(operationState=states.RUNNING_STATE), + Mock(operationState=states.FINISHED_STATE), ] cursor.fetch_logs.side_effect = timeout with ( app.app_context(), patch("superset.db_engine_specs.hive.db") as db, patch("superset.db_engine_specs.hive.time.sleep"), + patch.dict("sys.modules", {"pyhive": Mock(hive=hive), "pyhive.hive": hive}), ): db.session.query.return_value.filter_by.return_value.one.return_value = query query.status = "running" From 85ef2d19eb32be18d0531806302f5312320ce037 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Thu, 29 Oct 2026 22:16:16 +0000 Subject: [PATCH 03/14] test(impala): cover time limits during fetch cancellation --- superset/db_engine_specs/base.py | 9 ++--- .../unit_tests/db_engine_specs/test_impala.py | 34 +++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index a6e4c8a9eb18..1d5a93436591 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -136,10 +136,11 @@ # that we know this is a necessary evil we create an alias, and catch it instead. GenericDBException = Exception -# SQL Lab's time limits are raised into whatever the worker is executing: the -# Celery soft limit for asynchronous runs and the ``SQLLAB_TIMEOUT`` alarm for -# synchronous ones. Each fires once, so a handler that swallows it leaves the -# query running with nothing left to stop it. +# 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) diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 4020cc7f308a..c48343d3483a 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -333,3 +333,37 @@ def test_fetch_data_cancels_the_operation_on_time_limit( spec.fetch_data(cursor) cursor.cancel_operation.assert_called_once_with() cursor.fetchall.assert_not_called() + + +@pytest.mark.parametrize( + "cancel_timeout_type", [SoftTimeLimitExceeded, SupersetTimeoutException] +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_fetch_data_propagates_time_limit_during_cancel( + sleep_mock: Mock, + cancel_timeout_type: type[Exception], + app_context: None, + impala_errors: None, +) -> None: + """A second time limit in cancellation must replace the original time limit.""" + original_timeout = SoftTimeLimitExceeded("waiting") + cancel_timeout: Exception + if cancel_timeout_type is SupersetTimeoutException: + cancel_timeout = SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out during cancellation", + level=ErrorLevel.ERROR, + ) + else: + cancel_timeout = cancel_timeout_type("cancelling") + sleep_mock.side_effect = original_timeout + cursor = Mock() + cursor.is_executing.return_value = True + cursor.cancel_operation.side_effect = cancel_timeout + + with pytest.raises(cancel_timeout_type) as raised: + spec.fetch_data(cursor) + + assert raised.value is cancel_timeout + cursor.cancel_operation.assert_called_once_with() + cursor.fetchall.assert_not_called() From 30305b1dee182da078b4af460f4afb16d5ba3736 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Fri, 30 Oct 2026 00:16:58 +0000 Subject: [PATCH 04/14] fix(impala): cancel operations on row-fetch time limits --- .../configuration/configuring-superset.mdx | 4 +- superset/db_engine_specs/impala.py | 49 ++++++++++--------- .../unit_tests/db_engine_specs/test_impala.py | 37 ++++++++++++++ 3 files changed, 66 insertions(+), 24 deletions(-) diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 97cfb3f343e7..d6d66d269c4a 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -833,7 +833,9 @@ 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 Impala rows are being fetched also +request cancellation before the timeout propagates, including after execution has +finished. A failed cancellation RPC is logged without hiding the original timeout. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 7102402347e5..d7760a407b00 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -139,33 +139,36 @@ def execute( @classmethod def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]: """Wait for asynchronous operations using the public cursor API.""" - if callable(getattr(cursor, "execute_async", None)): - from impala.error import Error + try: + if callable(getattr(cursor, "execute_async", None)): + from impala.error import Error - deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"] + deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"] + try: + while cursor.is_executing(): + if time.monotonic() >= deadline: + cursor.cancel_operation() + raise TimeoutError( + "Timed out waiting for the Impala operation" + ) + time.sleep(0.1) + if cursor.execution_failed(): + # Fetching surfaces the driver's detailed asynchronous error, + # including errors on statements that produce no result set. + cursor.fetchall() + except Error as ex: + raise cls.get_dbapi_mapped_exception(ex) from ex + return super().fetch_data(cursor, limit) + except TIME_LIMIT_EXCEPTIONS: + # The operation is released before the time limit propagates so + # it does not keep running on the coordinator. try: - while cursor.is_executing(): - if time.monotonic() >= deadline: - cursor.cancel_operation() - raise TimeoutError("Timed out waiting for the Impala operation") - time.sleep(0.1) - if cursor.execution_failed(): - # Fetching surfaces the driver's detailed asynchronous error, - # including errors on statements that produce no result set. - cursor.fetchall() - except Error as ex: - raise cls.get_dbapi_mapped_exception(ex) from ex + cursor.cancel_operation() except TIME_LIMIT_EXCEPTIONS: - # The operation is released before the time limit propagates so - # it does not keep running on the coordinator. - try: - cursor.cancel_operation() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - logger.warning("cancel_operation() failed after a time limit") raise - return super().fetch_data(cursor, limit) + except Exception: # pylint: disable=broad-except + logger.warning("cancel_operation() failed after a time limit") + raise @classmethod def _cancel_operation(cls, cursor: Any, query_id: int) -> None: diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index c48343d3483a..a395da208a11 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -367,3 +367,40 @@ def test_fetch_data_propagates_time_limit_during_cancel( assert raised.value is cancel_timeout cursor.cancel_operation.assert_called_once_with() cursor.fetchall.assert_not_called() + + +@pytest.mark.parametrize("cancel_fails", [False, True]) +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out fetching rows", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +def test_fetch_data_cancels_on_time_limit_during_row_fetch( + timeout: Exception, + cancel_fails: bool, + app_context: None, + impala_errors: None, +) -> None: + """A time limit after polling ends cancels even when the cancel RPC fails.""" + from superset.db_engine_specs.base import BaseEngineSpec + + cursor = Mock() + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + if cancel_fails: + cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") + + with patch.object(BaseEngineSpec, "fetch_data", side_effect=timeout) as fetch: + with pytest.raises(type(timeout)) as raised: + spec.fetch_data(cursor, 10) + + assert raised.value is timeout + fetch.assert_called_once_with(cursor, 10) + cursor.cancel_operation.assert_called_once_with() From d66573316f59337fe1282d0e14d4d28349c67a70 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Fri, 30 Oct 2026 12:34:36 +0000 Subject: [PATCH 05/14] fix(impala): release deadline operations when cancellation fails --- superset/db_engine_specs/impala.py | 4 ++-- .../unit_tests/db_engine_specs/test_impala.py | 23 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index d7760a407b00..a08ce7f5b0ce 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -147,7 +147,7 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. try: while cursor.is_executing(): if time.monotonic() >= deadline: - cursor.cancel_operation() + cls._cancel_operation(cursor) raise TimeoutError( "Timed out waiting for the Impala operation" ) @@ -171,7 +171,7 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. raise @classmethod - def _cancel_operation(cls, cursor: Any, query_id: int) -> None: + def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: """Cancel the live operation and release its handles.""" try: cursor.cancel_operation() diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index a395da208a11..1cd6d5fc4e52 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -257,6 +257,29 @@ def test_fetch_data_timeout( assert monotonic_mock.call_count == 3 +@patch("superset.db_engine_specs.impala.time.monotonic", side_effect=[0, 31]) +def test_fetch_data_deadline_releases_operation_when_cancel_fails( + monotonic_mock: Mock, app_context: None, impala_errors: None +) -> None: + """A failed cancel RPC must not mask the deadline or retain operation handles.""" + from flask import current_app + from impala.error import Error + + cursor = Mock() + cursor.is_executing.return_value = True + cursor.cancel_operation.side_effect = Error("cancel RPC failed") + + with patch.dict(current_app.config, SQLLAB_TIMEOUT=30): + with pytest.raises(TimeoutError, match="Timed out waiting"): + spec.fetch_data(cursor) + + cursor.cancel_operation.assert_called_once_with() + cursor.close_operation.assert_called_once_with() + cursor.close.assert_called_once_with() + cursor.fetchall.assert_not_called() + assert monotonic_mock.call_count == 2 + + @pytest.mark.parametrize("error", [TypeError("bug"), AttributeError("bug")]) @patch.object(spec, "get_dbapi_mapped_exception") def test_fetch_data_preserves_programming_errors( From bde3fa4522a0590de2e3030760aa05a2d4064345 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Fri, 30 Oct 2026 23:18:34 +0000 Subject: [PATCH 06/14] fix(impala): reconnect before cancelling interrupted result fetch --- .../configuration/configuring-superset.mdx | 4 ++- superset/db_engine_specs/impala.py | 10 +++++-- .../unit_tests/db_engine_specs/test_impala.py | 29 ++++++++++++++++++- 3 files changed, 38 insertions(+), 5 deletions(-) diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index d6d66d269c4a..4b2156c3eabd 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -835,7 +835,9 @@ pending when result fetching begins, the additional polling wait is bounded by a timeout. Configure the driver's network timeout separately to bound individual RPC calls. SQL Lab time limits that fire while Impala rows are being fetched also request cancellation before the timeout propagates, including after execution has -finished. A failed cancellation RPC is logged without hiding the original timeout. +finished. Before cancellation, Superset reopens the driver transport to discard any +partially read result response, then releases the operation and cursor handles. +A failed cancellation RPC is logged without hiding the original timeout. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index a08ce7f5b0ce..a457bccd9fd9 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -161,13 +161,17 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. return super().fetch_data(cursor, limit) except TIME_LIMIT_EXCEPTIONS: # The operation is released before the time limit propagates so - # it does not keep running on the coordinator. + # it does not keep running on the coordinator. The limit can land + # part way through a FetchResults response, so the transport is + # reopened first: reading the cancel's reply from the remainder of + # that response would hang or request a garbage-sized read. try: - cursor.cancel_operation() + cursor.session.service.reconnect() except TIME_LIMIT_EXCEPTIONS: raise except Exception: # pylint: disable=broad-except - logger.warning("cancel_operation() failed after a time limit") + logger.warning("reconnect() failed after a time limit") + cls._cancel_operation(cursor) raise @classmethod diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 1cd6d5fc4e52..212bb7a9b7dd 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -17,7 +17,7 @@ from datetime import datetime from typing import Optional -from unittest.mock import Mock, patch +from unittest.mock import call, MagicMock, Mock, patch import pytest from celery.exceptions import SoftTimeLimitExceeded @@ -427,3 +427,30 @@ def test_fetch_data_cancels_on_time_limit_during_row_fetch( assert raised.value is timeout fetch.assert_called_once_with(cursor, 10) cursor.cancel_operation.assert_called_once_with() + + +@pytest.mark.parametrize("cancel_fails", [False, True]) +def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( + cancel_fails: bool, app_context: None, impala_errors: None +) -> None: + """Discard an interrupted response before cancellation and release handles.""" + from superset.db_engine_specs.base import BaseEngineSpec + + cursor = Mock() + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + calls = MagicMock() + calls.attach_mock(cursor.session.service.reconnect, "reconnect") + calls.attach_mock(cursor.cancel_operation, "cancel_operation") + if cancel_fails: + cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") + timeout = SoftTimeLimitExceeded() + + with patch.object(BaseEngineSpec, "fetch_data", side_effect=timeout): + with pytest.raises(SoftTimeLimitExceeded) as raised: + spec.fetch_data(cursor) + + assert raised.value is timeout + assert calls.mock_calls[:2] == [call.reconnect(), call.cancel_operation()] + cursor.close_operation.assert_called_once_with() + cursor.close.assert_called_once_with() From 24860720185c06394e8eb4ccf5c0f4a30cea21b2 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 01:22:05 +0000 Subject: [PATCH 07/14] fix(impala): reconnect before cancelling interrupted polling --- .../configuration/configuring-superset.mdx | 9 ++-- superset/db_engine_specs/impala.py | 31 +++++++------ .../unit_tests/db_engine_specs/test_impala.py | 46 +++++++++++++++++++ 3 files changed, 69 insertions(+), 17 deletions(-) diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 4b2156c3eabd..efa8470561df 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -833,10 +833,11 @@ 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. SQL Lab time limits that fire while Impala rows are being fetched also -request cancellation before the timeout propagates, including after execution has -finished. Before cancellation, Superset reopens the driver transport to discard any -partially read result response, then releases the operation and cursor handles. +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. Before cancellation, Superset reopens the +driver transport to discard any partially read RPC response, then releases the +operation and cursor handles. A failed cancellation RPC is logged without hiding the original timeout. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index a457bccd9fd9..151184e3ab46 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -160,20 +160,25 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. raise cls.get_dbapi_mapped_exception(ex) from ex return super().fetch_data(cursor, limit) except TIME_LIMIT_EXCEPTIONS: - # The operation is released before the time limit propagates so - # it does not keep running on the coordinator. The limit can land - # part way through a FetchResults response, so the transport is - # reopened first: reading the cancel's reply from the remainder of - # that response would hang or request a garbage-sized read. - try: - cursor.session.service.reconnect() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - logger.warning("reconnect() failed after a time limit") - cls._cancel_operation(cursor) + cls._cancel_operation_after_time_limit(cursor) raise + @classmethod + def _cancel_operation_after_time_limit( + cls, cursor: Any, query_id: int | None = None + ) -> None: + """Discard interrupted RPC replies before cancelling and releasing handles.""" + # A signal can interrupt any RPC reply, including status, logs, or rows. + # Reading the cancel reply on that transport could consume leftover bytes + # and hang or request a garbage-sized read, so reopen the transport first. + try: + cursor.session.service.reconnect() + except TIME_LIMIT_EXCEPTIONS: + raise + except Exception: # pylint: disable=broad-except + logger.warning("reconnect() failed after a time limit") + cls._cancel_operation(cursor, query_id) + @classmethod def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: """Cancel the live operation and release its handles.""" @@ -260,7 +265,7 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: # SQL Lab ends the query as timed out once this propagates; the # operation is cancelled first so it does not keep running on the # coordinator while fetch_data() waits on it. - cls._cancel_operation(cursor, query_id) + cls._cancel_operation_after_time_limit(cursor, query_id) raise except Exception: # pylint: disable=broad-except logger.debug("Call to status() failed ") diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 212bb7a9b7dd..4411f2d8843c 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -454,3 +454,49 @@ def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( assert calls.mock_calls[:2] == [call.reconnect(), call.cancel_operation()] cursor.close_operation.assert_called_once_with() cursor.close.assert_called_once_with() + + +@pytest.mark.parametrize("interrupted_rpc", ["status", "get_log"]) +@pytest.mark.parametrize("cancel_fails", [False, True]) +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out polling the operation", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +def test_handle_cursor_reconnects_before_cancelling_an_interrupted_poll( + interrupted_rpc: str, + cancel_fails: bool, + timeout: Exception, + app_context: None, +) -> None: + """Discard interrupted status/log replies before cancelling and releasing.""" + cursor = Mock() + cursor.status.return_value = "RUNNING_STATE" + getattr(cursor, interrupted_rpc).side_effect = timeout + if cancel_fails: + cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") + query = Mock(id=1, extra={}, status="running", progress=0) + calls = MagicMock() + calls.attach_mock(cursor.session.service.reconnect, "reconnect") + calls.attach_mock(cursor.cancel_operation, "cancel_operation") + calls.attach_mock(cursor.close_operation, "close_operation") + calls.attach_mock(cursor.close, "close") + + with patch("superset.db_engine_specs.impala.db"): + with pytest.raises(type(timeout)) as raised: + spec.handle_cursor(cursor, query) + + assert raised.value is timeout + assert calls.mock_calls == [ + call.reconnect(), + call.cancel_operation(), + call.close_operation(), + call.close(), + ] From ca8cbd0680a8a1d5a3748aa8759737cb34b0b5ee Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 04:30:07 +0000 Subject: [PATCH 08/14] fix(impala): empty transport buffers before cancelling after a time limit impyla's reconnect() only closes and reopens the socket or HTTP connection. TBufferedTransport and TSaslClientTransport keep their read and write buffers across close and open, so the unread rest of an interrupted status, log or row reply was parsed as the CancelOperation reply: on the binary and SASL transports the cancel blocked until the socket timeout and impyla resent it, and over HTTP it failed outright. Reset the transport instead: check every layer of the stack, close it, replace each layer's buffers with empty ones and reopen it. If the stack is not recognised or the reset fails, no cancel RPC is sent on it. Tests drive impyla over real binary, SASL PLAIN and HTTP transports against a local HiveServer2 stand-in, with the time limit raised right after a reply header is read. --- .../configuration/configuring-superset.mdx | 13 +- superset/db_engine_specs/impala.py | 92 ++++- .../unit_tests/db_engine_specs/test_impala.py | 91 +++- .../db_engine_specs/test_impala_pending.py | 15 + .../db_engine_specs/test_impala_transport.py | 391 ++++++++++++++++++ 5 files changed, 585 insertions(+), 17 deletions(-) create mode 100644 tests/unit_tests/db_engine_specs/test_impala_transport.py diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index efa8470561df..21b8e7871e97 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -835,10 +835,15 @@ pending when result fetching begins, the additional polling wait is bounded by a timeout. Configure the driver's network timeout separately to bound individual 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. Before cancellation, Superset reopens the -driver transport to discard any partially read RPC response, then releases the -operation and cursor handles. -A failed cancellation RPC is logged without hiding the original timeout. +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. If the transport cannot be reset, no cancellation RPC is sent and a +warning is logged; the operation is then left to Impala's own query and session +timeouts. A failed cancellation RPC is logged without hiding the original +timeout. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 151184e3ab46..06c9f990b7ec 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -167,18 +167,100 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. def _cancel_operation_after_time_limit( cls, cursor: Any, query_id: int | None = None ) -> None: - """Discard interrupted RPC replies before cancelling and releasing handles.""" + """Cancel and release the operation on a transport with no stale bytes.""" # A signal can interrupt any RPC reply, including status, logs, or rows. - # Reading the cancel reply on that transport could consume leftover bytes - # and hang or request a garbage-sized read, so reopen the transport first. + # The unread rest of that reply stays in the socket and in the + # transport's read buffer, which survives a plain close and reopen, so + # the cancel reply would be parsed from it and could hang or request a + # garbage-sized read. Without a transport known to be clean, no RPC is + # sent and Impala's own timeouts close the operation. try: - cursor.session.service.reconnect() + reset = cls._reset_transport(cursor) except TIME_LIMIT_EXCEPTIONS: raise except Exception: # pylint: disable=broad-except - logger.warning("reconnect() failed after a time limit") + logger.warning("Query %s: resetting the transport failed", query_id) + reset = False + if not reset: + logger.warning( + "Query %s: not cancelling the operation on an unclean transport", + query_id, + ) + return cls._cancel_operation(cursor, query_id) + @staticmethod + def _reset_transport(cursor: Any) -> bool: + """ + Reopen the cursor's Thrift transport with every buffer emptied. + + impyla's ``reconnect()`` only closes and reopens the socket or HTTP + connection, while the buffering layers above it keep their read and + write buffers. Returns whether every layer was recognised and reset. + """ + # pylint: disable=import-outside-toplevel + from io import BytesIO + + from impala._thrift_api import ImpalaHttpClient + from thrift.transport.TSocket import TSocket + from thrift.transport.TTransport import TBufferedTransport + + # (layer type, its buffer attributes, the attribute of the layer below). + # The buffers are private to thrift, thrift_sasl and impyla, so the + # whole stack is checked before anything is changed. + layouts: list[tuple[type, tuple[str, ...], str | None]] = [ + ( + TBufferedTransport, + ("_TBufferedTransport__rbuf", "_TBufferedTransport__wbuf"), + "_TBufferedTransport__trans", + ), + # Each HTTP request opens a new connection; only the pending + # request body is kept on the client. + (ImpalaHttpClient, ("_ImpalaHttpClient__wbuf",), None), + ] + try: + from thrift_sasl import TSaslClientTransport + + layouts.append( + ( + TSaslClientTransport, + ("_TSaslClientTransport__rbuf", "_TSaslClientTransport__wbuf"), + "_trans", + ) + ) + except ImportError: + pass + + # The same transport object impyla's own ``reconnect()`` reopens. + transport = cursor.session.service.client._iprot.trans + buffers: list[tuple[Any, str]] = [] + layer = transport + while not isinstance(layer, TSocket): + layout = next((lay for lay in layouts if isinstance(layer, lay[0])), None) + if layout is None: + return False + _, attributes, inner = layout + if not all(hasattr(layer, attribute) for attribute in attributes): + return False + buffers.extend((layer, attribute) for attribute in attributes) + if inner is None: + break + layer = getattr(layer, inner, None) + + # Closing drops the socket together with any reply bytes still in it. + try: + transport.close() + except TIME_LIMIT_EXCEPTIONS: + raise + except Exception: # pylint: disable=broad-except + # An HTTP client that was already closed raises on a second close. + if transport.isOpen(): + raise + for buffered_layer, attribute in buffers: + setattr(buffered_layer, attribute, BytesIO()) + transport.open() + return True + @classmethod def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: """Cancel the live operation and release its handles.""" diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 4411f2d8843c..e70984aeef58 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Iterator from datetime import datetime -from typing import Optional +from typing import Any, Optional from unittest.mock import call, MagicMock, Mock, patch import pytest @@ -185,6 +186,18 @@ class _DriverError(Exception): """Represent impyla's optional DB-API error base in unit tests.""" +@pytest.fixture(autouse=True) +def clean_transport() -> Iterator[Mock]: + """ + Report the transport reset as successful for mock cursors. + + Mock cursors have no Thrift transport; resets over real impyla transports + are covered in ``test_impala_transport.py``. + """ + with patch.object(spec, "_reset_transport", return_value=True) as reset_transport: + yield reset_transport + + @pytest.fixture def impala_errors(mocker: MockerFixture) -> None: """Provide driver errors without requiring the optional impyla package.""" @@ -430,8 +443,11 @@ def test_fetch_data_cancels_on_time_limit_during_row_fetch( @pytest.mark.parametrize("cancel_fails", [False, True]) -def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( - cancel_fails: bool, app_context: None, impala_errors: None +def test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch( + cancel_fails: bool, + clean_transport: Mock, + app_context: None, + impala_errors: None, ) -> None: """Discard an interrupted response before cancellation and release handles.""" from superset.db_engine_specs.base import BaseEngineSpec @@ -440,7 +456,7 @@ def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( cursor.is_executing.return_value = False cursor.execution_failed.return_value = False calls = MagicMock() - calls.attach_mock(cursor.session.service.reconnect, "reconnect") + calls.attach_mock(clean_transport, "reset_transport") calls.attach_mock(cursor.cancel_operation, "cancel_operation") if cancel_fails: cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") @@ -451,7 +467,10 @@ def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( spec.fetch_data(cursor) assert raised.value is timeout - assert calls.mock_calls[:2] == [call.reconnect(), call.cancel_operation()] + assert calls.mock_calls[:2] == [ + call.reset_transport(cursor), + call.cancel_operation(), + ] cursor.close_operation.assert_called_once_with() cursor.close.assert_called_once_with() @@ -470,10 +489,11 @@ def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch( ], ids=["async", "sync"], ) -def test_handle_cursor_reconnects_before_cancelling_an_interrupted_poll( +def test_handle_cursor_resets_transport_before_cancelling_an_interrupted_poll( interrupted_rpc: str, cancel_fails: bool, timeout: Exception, + clean_transport: Mock, app_context: None, ) -> None: """Discard interrupted status/log replies before cancelling and releasing.""" @@ -484,7 +504,7 @@ def test_handle_cursor_reconnects_before_cancelling_an_interrupted_poll( cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") query = Mock(id=1, extra={}, status="running", progress=0) calls = MagicMock() - calls.attach_mock(cursor.session.service.reconnect, "reconnect") + calls.attach_mock(clean_transport, "reset_transport") calls.attach_mock(cursor.cancel_operation, "cancel_operation") calls.attach_mock(cursor.close_operation, "close_operation") calls.attach_mock(cursor.close, "close") @@ -495,8 +515,63 @@ def test_handle_cursor_reconnects_before_cancelling_an_interrupted_poll( assert raised.value is timeout assert calls.mock_calls == [ - call.reconnect(), + call.reset_transport(cursor), call.cancel_operation(), call.close_operation(), call.close(), ] + + +@pytest.mark.parametrize( + "reset_outcome", + [{"return_value": False}, {"side_effect": OSError("connection refused")}], + ids=["unknown_transport", "reopen_fails"], +) +@pytest.mark.parametrize("interrupted_rpc", ["status", "get_log", "fetch"]) +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out polling the operation", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +def test_time_limit_skips_cancel_rpcs_without_a_clean_transport( + reset_outcome: dict[str, Any], + interrupted_rpc: str, + timeout: Exception, + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """No RPC is sent on a transport that may hold an interrupted reply.""" + from superset.db_engine_specs.base import BaseEngineSpec + + clean_transport.configure_mock(**reset_outcome) + cursor = Mock() + cursor.status.return_value = "RUNNING_STATE" + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + query = Mock(id=1, extra={}, status="running", progress=0) + + with ( + patch("superset.db_engine_specs.impala.db"), + patch.object(BaseEngineSpec, "fetch_data", side_effect=timeout), + ): + if interrupted_rpc == "fetch": + with pytest.raises(type(timeout)) as raised: + spec.fetch_data(cursor) + else: + getattr(cursor, interrupted_rpc).side_effect = timeout + with pytest.raises(type(timeout)) as raised: + spec.handle_cursor(cursor, query) + + assert raised.value is timeout + clean_transport.assert_called_once_with(cursor) + cursor.cancel_operation.assert_not_called() + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 0ad054a71d3f..7110d92d012b 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_pending.py +++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Iterator from contextlib import nullcontext from unittest.mock import MagicMock, Mock, patch @@ -41,6 +42,20 @@ from superset.utils.core import QueryStatus +@pytest.fixture(autouse=True) +def clean_transport() -> Iterator[Mock]: + """ + Report the transport reset as successful for mock cursors. + + Mock cursors have no Thrift transport; resets over real impyla transports + are covered in ``test_impala_transport.py``. + """ + with patch.object( + ImpalaEngineSpec, "_reset_transport", return_value=True + ) as reset_transport: + yield reset_transport + + @pytest.mark.parametrize("use_executor", [False, True]) def test_stop_during_poll_does_not_fail_sync_fetch( mocker: MockerFixture, app_context: None, use_executor: bool diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py new file mode 100644 index 000000000000..7edfcc96d141 --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -0,0 +1,391 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Time-limit cancellation over real impyla Thrift transports. + +A local HiveServer2 stand-in serves impyla's binary, SASL PLAIN and HTTP +transports. The time limit is raised right after a reply header is read, as a +signal landing at that point would, so the rest of the reply is left in the +client's buffers. +""" + +from __future__ import annotations + +import logging +import socket +import struct +import threading +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from celery.exceptions import SoftTimeLimitExceeded + +from superset.db_engine_specs.impala import ImpalaEngineSpec as spec # noqa: N813 +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import SupersetTimeoutException + +hs2 = pytest.importorskip("impala.hiveserver2") +ttypes = pytest.importorskip("impala._thrift_gen.TCLIService.ttypes") +TCLIService = pytest.importorskip("impala._thrift_gen.TCLIService.TCLIService") +TSocketModule = pytest.importorskip("thrift.transport.TSocket") +TTransportModule = pytest.importorskip("thrift.transport.TTransport") +TBinaryProtocol = pytest.importorskip("thrift.protocol.TBinaryProtocol") +THttpServer = pytest.importorskip("thrift.server.THttpServer") + +OK = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) +OPERATION_GUID = b"superset-op-guid" + + +@dataclass +class FakeHiveServer2: + """The HiveServer2 RPCs impyla issues for one SQL Lab query.""" + + state: int = ttypes.TOperationState.RUNNING_STATE + calls: list[tuple[str, int, Any]] = field(default_factory=list) + connection: int = 0 + + def _record(self, method: str, req: Any) -> None: + self.calls.append((method, self.connection, req)) + + def OpenSession(self, req: Any) -> Any: # noqa: N802 + self._record("OpenSession", req) + return ttypes.TOpenSessionResp( + status=OK, + serverProtocolVersion=ttypes.TProtocolVersion.HIVE_CLI_SERVICE_PROTOCOL_V6, + sessionHandle=ttypes.TSessionHandle( + sessionId=ttypes.THandleIdentifier(guid=b"s" * 16, secret=b"s" * 16) + ), + configuration={}, + ) + + def ExecuteStatement(self, req: Any) -> Any: # noqa: N802 + self._record("ExecuteStatement", req) + return ttypes.TExecuteStatementResp( + status=OK, + operationHandle=ttypes.TOperationHandle( + operationId=ttypes.THandleIdentifier( + guid=OPERATION_GUID, secret=b"o" * 16 + ), + operationType=ttypes.TOperationType.EXECUTE_STATEMENT, + hasResultSet=True, + ), + ) + + def GetOperationStatus(self, req: Any) -> Any: # noqa: N802 + self._record("GetOperationStatus", req) + return ttypes.TGetOperationStatusResp( + status=OK, operationState=self.state, hasResultSet=True + ) + + def GetLog(self, req: Any) -> Any: # noqa: N802 + self._record("GetLog", req) + # Longer than one 4 KiB transport buffer. + return ttypes.TGetLogResp(status=OK, log="Query q: 5% Complete\n" * 1000) + + def GetResultSetMetadata(self, req: Any) -> Any: # noqa: N802 + self._record("GetResultSetMetadata", req) + string_type = ttypes.TTypeDesc( + types=[ + ttypes.TTypeEntry( + primitiveEntry=ttypes.TPrimitiveTypeEntry( + type=ttypes.TTypeId.STRING_TYPE + ) + ) + ] + ) + return ttypes.TGetResultSetMetadataResp( + status=OK, + schema=ttypes.TTableSchema( + columns=[ + ttypes.TColumnDesc(columnName="c", typeDesc=string_type, position=1) + ] + ), + ) + + def FetchResults(self, req: Any) -> Any: # noqa: N802 + self._record("FetchResults", req) + values = [b"row value " * 5] * 500 + return ttypes.TFetchResultsResp( + status=OK, + hasMoreRows=False, + results=ttypes.TRowSet( + startRowOffset=0, + rows=[], + columns=[ + ttypes.TColumn( + stringVal=ttypes.TStringColumn( + values=values, nulls=bytes(len(values) // 8 + 1) + ) + ) + ], + ), + ) + + def CancelOperation(self, req: Any) -> Any: # noqa: N802 + self._record("CancelOperation", req) + return ttypes.TCancelOperationResp(status=OK) + + def CloseOperation(self, req: Any) -> Any: # noqa: N802 + self._record("CloseOperation", req) + return ttypes.TCloseOperationResp(status=OK) + + def CloseSession(self, req: Any) -> Any: # noqa: N802 + self._record("CloseSession", req) + return ttypes.TCloseSessionResp(status=OK) + + +def _read_exactly(conn: socket.socket, size: int) -> bytes: + data = b"" + while len(data) < size: + chunk = conn.recv(size - len(data)) + if not chunk: + raise EOFError + data += chunk + return data + + +def _sasl_plain_handshake(conn: socket.socket) -> None: + """Accept thrift_sasl's PLAIN negotiation: START, OK, then COMPLETE.""" + for _ in range(2): + _, length = struct.unpack(">BI", _read_exactly(conn, 5)) + _read_exactly(conn, length) + conn.sendall(struct.pack(">BI", 5, 0)) + + +def _serve_socket( + handler: FakeHiveServer2, listener: socket.socket, sasl: bool +) -> None: + processor = TCLIService.Processor(handler) + while True: + try: + conn, _ = listener.accept() + except OSError: + return + handler.connection += 1 + sock = TSocketModule.TSocket() + sock.handle = conn + try: + if sasl: + _sasl_plain_handshake(conn) + # PLAIN without a security layer frames each message with + # its length. + transport = TTransportModule.TFramedTransport(sock) + else: + transport = TTransportModule.TBufferedTransport(sock) + protocol = TBinaryProtocol.TBinaryProtocol(transport) + while True: + processor.process(protocol, protocol) + except Exception: # noqa: S110 # pylint: disable=broad-except + pass + finally: + conn.close() + + +@pytest.fixture(params=["binary", "sasl", "http"]) +def hs2_connection( + request: pytest.FixtureRequest, +) -> Iterator[tuple[str, FakeHiveServer2, Any]]: + """An impyla connection to a local HiveServer2 over each transport.""" + from impala.dbapi import connect + + mode = request.param + handler = FakeHiveServer2() + if mode == "http": + server = THttpServer.THttpServer( + TCLIService.Processor(handler), + ("127.0.0.1", 0), + TBinaryProtocol.TBinaryProtocolFactory(), + ) + server.httpd.timeout = 2 + port = server.httpd.server_address[1] + thread = threading.Thread(target=server.serve, daemon=True) + thread.start() + conn = connect( + "127.0.0.1", + port, + use_http_transport=True, + http_path="cliservice", + auth_mechanism="NOSASL", + ) + # The HTTP transport ignores ``connect(timeout=...)``. + conn.service.client._iprot.trans._TBufferedTransport__trans.setTimeout(2000) + try: + yield mode, handler, conn + finally: + server.httpd.shutdown() + server.httpd.server_close() + return + + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen() + thread = threading.Thread( + target=_serve_socket, args=(handler, listener, mode == "sasl"), daemon=True + ) + thread.start() + conn = connect( + "127.0.0.1", + listener.getsockname()[1], + timeout=2, + auth_mechanism="PLAIN" if mode == "sasl" else "NOSASL", + user="user", + password="password", # noqa: S106 + ) + try: + yield mode, handler, conn + finally: + listener.close() + + +def _interrupt_after_header(cursor: Any, reply: str, timeout: Exception) -> None: + """Raise ``timeout`` once the header of the first ``reply`` has been read.""" + protocol = cursor.session.service.client._iprot + read_message_begin = protocol.readMessageBegin + armed = [True] + + def interrupted_read_message_begin() -> Any: + header = read_message_begin() + if armed[0] and header[0] == reply: + armed[0] = False + raise timeout + return header + + protocol.readMessageBegin = interrupted_read_message_begin + + +TIMEOUTS = [ + pytest.param(SoftTimeLimitExceeded(), id="async"), + pytest.param( + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + id="sync", + ), +] + + +def _assert_cancelled_cleanly( + mode: str, + handler: FakeHiveServer2, + interrupted: str, + caplog: pytest.LogCaptureFixture, +) -> None: + methods = [method for method, _, _ in handler.calls] + # A cancel that read leftover bytes as its reply blocks until the socket + # timeout, after which impyla sends it again on a reopened transport. + assert methods.count("CancelOperation") == 1, methods + cancel_index = methods.index("CancelOperation") + _, cancel_connection, cancel = handler.calls[cancel_index] + assert cancel.operationHandle.operationId.guid == OPERATION_GUID + assert "CloseOperation" in methods[cancel_index:], methods + interrupted_connection = max( + connection + for method, connection, _ in handler.calls[:cancel_index] + if method == interrupted + ) + if mode != "http": + # The interrupted reply's connection is dropped, not reused. + assert cancel_connection > interrupted_connection + # Every cancellation RPC read its own reply. + failures = [ + record.getMessage() + for record in caplog.records + if record.name == "superset.db_engine_specs.impala" + and record.levelno >= logging.WARNING + ] + assert failures == [] + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "interrupted, state", + [ + ("GetOperationStatus", ttypes.TOperationState.RUNNING_STATE), + ("GetLog", ttypes.TOperationState.RUNNING_STATE), + ], + ids=["status", "get_log"], +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_handle_cursor_cancels_on_a_clean_transport_after_interrupted_reply( + sleep_mock: Mock, + interrupted: str, + state: int, + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, + caplog: pytest.LogCaptureFixture, +) -> None: + """Leftover reply bytes must not be read as the CancelOperation reply.""" + mode, handler, conn = hs2_connection + handler.state = state + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + _interrupt_after_header(cursor, interrupted, timeout) + query = Mock(id=1, extra={}, status="running", progress=0) + + with patch("superset.db_engine_specs.impala.db"): + with pytest.raises(type(timeout)) as raised: + spec.handle_cursor(cursor, query) + + assert raised.value is timeout + _assert_cancelled_cleanly(mode, handler, interrupted, caplog) + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "interrupted, state", + [ + ("GetOperationStatus", ttypes.TOperationState.RUNNING_STATE), + ("FetchResults", ttypes.TOperationState.FINISHED_STATE), + ], + ids=["status", "rows"], +) +def test_fetch_data_cancels_on_a_clean_transport_after_interrupted_reply( + interrupted: str, + state: int, + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, + caplog: pytest.LogCaptureFixture, +) -> None: + """Leftover status or row bytes must not be read as the cancel reply.""" + mode, handler, conn = hs2_connection + handler.state = state + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + _interrupt_after_header(cursor, interrupted, timeout) + + with pytest.raises(type(timeout)) as raised: + spec.fetch_data(cursor) + + assert raised.value is timeout + _assert_cancelled_cleanly(mode, handler, interrupted, caplog) + + +def test_reset_transport_refuses_an_unknown_transport() -> None: + """An unrecognised buffering layer is left untouched and reported unclean.""" + cursor = Mock() + transport = cursor.session.service.client._iprot.trans + + assert spec._reset_transport(cursor) is False + transport.close.assert_not_called() + transport.open.assert_not_called() From 365c5e877d6cb893da2e2b590bed770a4c2c4c41 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 04:38:19 +0000 Subject: [PATCH 09/14] fix(impala): send no cursor cleanup RPC after a skipped cancel When the transport reset is skipped or fails, the cancel RPC was not sent, but impyla's cursor close() (reached through closing() in the executor) and its finalizer still sent CloseOperation and CloseSession, reopening the transport with the interrupted reply still buffered. Mark the cursor closed on that path, and on a second time limit during the reset, so neither sends anything. --- .../configuration/configuring-superset.mdx | 6 +- superset/db_engine_specs/impala.py | 11 +++ .../unit_tests/db_engine_specs/test_impala.py | 40 +++++++++++ .../db_engine_specs/test_impala_transport.py | 68 ++++++++++++++++++- 4 files changed, 120 insertions(+), 5 deletions(-) diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 21b8e7871e97..b83e8e60c601 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -840,9 +840,9 @@ 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. If the transport cannot be reset, no cancellation RPC is sent and a -warning is logged; the operation is then left to Impala's own query and session -timeouts. A failed cancellation RPC is logged without hiding the original +handles. If the transport cannot be reset, a warning is logged and no +cancellation or cursor cleanup RPC is sent on it; the operation is then left to +Impala's own query and session timeouts. A failed cancellation RPC is logged without hiding the original timeout. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 06c9f990b7ec..3db6820e78ff 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -177,6 +177,7 @@ def _cancel_operation_after_time_limit( try: reset = cls._reset_transport(cursor) except TIME_LIMIT_EXCEPTIONS: + cls._abandon_cursor(cursor) raise except Exception: # pylint: disable=broad-except logger.warning("Query %s: resetting the transport failed", query_id) @@ -186,9 +187,19 @@ def _cancel_operation_after_time_limit( "Query %s: not cancelling the operation on an unclean transport", query_id, ) + cls._abandon_cursor(cursor) return cls._cancel_operation(cursor, query_id) + @staticmethod + def _abandon_cursor(cursor: Any) -> None: + """Keep impyla from sending RPCs for a cursor whose transport is unclean.""" + # ``HiveServer2Cursor.close()``, which callers reach through + # ``closing()``, and its finalizer send CloseOperation and CloseSession, + # reopening a closed transport with its buffers as they are. Both skip + # a cursor already marked closed. + cursor._closed = True # pylint: disable=protected-access + @staticmethod def _reset_transport(cursor: Any) -> bool: """ diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index e70984aeef58..5e7c469bc921 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -575,3 +575,43 @@ def test_time_limit_skips_cancel_rpcs_without_a_clean_transport( cursor.cancel_operation.assert_not_called() cursor.close_operation.assert_not_called() cursor.close.assert_not_called() + # impyla's close() and finalizer skip a cursor marked closed. + assert cursor._closed is True + + +@pytest.mark.parametrize( + "second_timeout_type", + [SoftTimeLimitExceeded, SupersetTimeoutException], +) +def test_time_limit_during_transport_reset_propagates( + second_timeout_type: type[Exception], + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """A second time limit while resetting propagates and sends no cancel RPC.""" + from superset.db_engine_specs.base import BaseEngineSpec + + second_timeout: Exception + if second_timeout_type is SoftTimeLimitExceeded: + second_timeout = SoftTimeLimitExceeded("resetting") + else: + second_timeout = SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out resetting the transport", + level=ErrorLevel.ERROR, + ) + clean_transport.side_effect = second_timeout + cursor = Mock() + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + + with patch.object( + BaseEngineSpec, "fetch_data", side_effect=SoftTimeLimitExceeded("fetching") + ): + with pytest.raises(type(second_timeout)) as raised: + spec.fetch_data(cursor) + + assert raised.value is second_timeout + cursor.cancel_operation.assert_not_called() + assert cursor._closed is True diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py index 7edfcc96d141..4f08696d0812 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_transport.py +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -30,6 +30,7 @@ import struct import threading from collections.abc import Iterator +from contextlib import closing from dataclasses import dataclass, field from typing import Any from unittest.mock import Mock, patch @@ -59,6 +60,8 @@ class FakeHiveServer2: state: int = ttypes.TOperationState.RUNNING_STATE calls: list[tuple[str, int, Any]] = field(default_factory=list) + # Every request received, including methods this stand-in does not serve. + requests: list[str] = field(default_factory=list) connection: int = 0 def _record(self, method: str, req: Any) -> None: @@ -151,6 +154,28 @@ def CloseSession(self, req: Any) -> Any: # noqa: N802 return ttypes.TCloseSessionResp(status=OK) +def _recording_protocol(transport: Any, handler: FakeHiveServer2) -> Any: + """Server-side protocol that records the name of every request.""" + protocol = TBinaryProtocol.TBinaryProtocol(transport) + read_message_begin = protocol.readMessageBegin + + def recording_read_message_begin() -> Any: + header = read_message_begin() + handler.requests.append(header[0]) + return header + + protocol.readMessageBegin = recording_read_message_begin + return protocol + + +class _RecordingProtocolFactory: + def __init__(self, handler: FakeHiveServer2) -> None: + self.handler = handler + + def getProtocol(self, transport: Any) -> Any: # noqa: N802 + return _recording_protocol(transport, self.handler) + + def _read_exactly(conn: socket.socket, size: int) -> bytes: data = b"" while len(data) < size: @@ -189,7 +214,7 @@ def _serve_socket( transport = TTransportModule.TFramedTransport(sock) else: transport = TTransportModule.TBufferedTransport(sock) - protocol = TBinaryProtocol.TBinaryProtocol(transport) + protocol = _recording_protocol(transport, handler) while True: processor.process(protocol, protocol) except Exception: # noqa: S110 # pylint: disable=broad-except @@ -211,7 +236,7 @@ def hs2_connection( server = THttpServer.THttpServer( TCLIService.Processor(handler), ("127.0.0.1", 0), - TBinaryProtocol.TBinaryProtocolFactory(), + _RecordingProtocolFactory(handler), ) server.httpd.timeout = 2 port = server.httpd.server_address[1] @@ -381,6 +406,45 @@ def test_fetch_data_cancels_on_a_clean_transport_after_interrupted_reply( _assert_cancelled_cleanly(mode, handler, interrupted, caplog) +@pytest.mark.parametrize( + "reset_failure", + [{"return_value": False}, {"side_effect": OSError("transport close failed")}], + ids=["unknown_transport", "reset_raises"], +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_closing_the_cursor_sends_nothing_after_a_skipped_cancel( + sleep_mock: Mock, + reset_failure: dict[str, Any], + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """ + Without a clean transport, cursor cleanup must not read leftover bytes. + + Callers close the cursor through ``closing()``; impyla's ``close()`` and + finalizer would otherwise send CloseOperation and CloseSession on the + transport holding the interrupted reply. + """ + _, handler, conn = hs2_connection + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + _interrupt_after_header(cursor, "GetOperationStatus", SoftTimeLimitExceeded()) + query = Mock(id=1, extra={}, status="running", progress=0) + + with ( + patch.object(spec, "_reset_transport", **reset_failure), + patch("superset.db_engine_specs.impala.db"), + ): + with pytest.raises(SoftTimeLimitExceeded): + with closing(cursor): + spec.handle_cursor(cursor, query) + cursor.__del__() + + # The first status poll is the interrupted RPC. + interrupted = handler.requests.index("GetOperationStatus") + assert handler.requests[interrupted + 1 :] == [] + + def test_reset_transport_refuses_an_unknown_transport() -> None: """An unrecognised buffering layer is left untouched and reported unclean.""" cursor = Mock() From 2e1b2559617d4e5089b4d88bf1405f1ed972a0bd Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 13:44:50 +0000 Subject: [PATCH 10/14] fix(impala): abandon interrupted cleanup on the real cursor --- .../configuration/configuring-superset.mdx | 12 +- superset/db_engine_specs/hive.py | 53 +++- superset/db_engine_specs/impala.py | 173 +++++++++---- superset/sql/execution/executor.py | 20 +- tests/unit_tests/db_engine_specs/test_hive.py | 49 ++++ .../db_engine_specs/test_hive_transport.py | 203 ++++++++++++++++ .../unit_tests/db_engine_specs/test_impala.py | 151 ++++++++++-- .../db_engine_specs/test_impala_pending.py | 51 +++- .../db_engine_specs/test_impala_transport.py | 227 +++++++++++++++++- .../unit_tests/sql/execution/test_executor.py | 22 ++ 10 files changed, 853 insertions(+), 108 deletions(-) create mode 100644 tests/unit_tests/db_engine_specs/test_hive_transport.py diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index b83e8e60c601..6e164309165a 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -840,10 +840,14 @@ 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. If the transport cannot be reset, a warning is logged and no -cancellation or cursor cleanup RPC is sent on it; the operation is then left to -Impala's own query and session timeouts. A failed cancellation RPC is logged without hiding the original -timeout. +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 diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index c1507f4f322b..40bb7bcd64d3 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -202,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( @@ -433,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 diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 3db6820e78ff..02a5fd8421c5 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -168,6 +168,10 @@ def _cancel_operation_after_time_limit( cls, cursor: Any, query_id: int | None = None ) -> None: """Cancel and release the operation on a transport with no stale bytes.""" + # A closed cursor has no operation left to cancel, and an abandoned one + # must send nothing more. + if getattr(cursor, "_closed", False) is True: + return # A signal can interrupt any RPC reply, including status, logs, or rows. # The unread rest of that reply stays in the socket and in the # transport's read buffer, which survives a plain close and reopen, so @@ -176,10 +180,9 @@ def _cancel_operation_after_time_limit( # sent and Impala's own timeouts close the operation. try: reset = cls._reset_transport(cursor) - except TIME_LIMIT_EXCEPTIONS: - cls._abandon_cursor(cursor) - raise except Exception: # pylint: disable=broad-except + # Cleanup must not replace the original time-limit exception, + # including when a second signal interrupts the reset itself. logger.warning("Query %s: resetting the transport failed", query_id) reset = False if not reset: @@ -187,31 +190,88 @@ def _cancel_operation_after_time_limit( "Query %s: not cancelling the operation on an unclean transport", query_id, ) - cls._abandon_cursor(cursor) + cls._abandon_cursor(cursor, query_id) return - cls._cancel_operation(cursor, query_id) + try: + cls._cancel_operation(cursor, query_id) + except TIME_LIMIT_EXCEPTIONS: + # Another limit landed inside a cleanup reply, so the transport + # holds unread bytes again. + logger.warning("Query %s: time limit during cancellation cleanup", query_id) - @staticmethod - def _abandon_cursor(cursor: Any) -> None: - """Keep impyla from sending RPCs for a cursor whose transport is unclean.""" + @classmethod + def _abandon_cursor(cls, cursor: Any, query_id: int | None = None) -> None: + """Send no further RPC on a transport that may hold an unread reply.""" # ``HiveServer2Cursor.close()``, which callers reach through - # ``closing()``, and its finalizer send CloseOperation and CloseSession, - # reopening a closed transport with its buffers as they are. Both skip - # a cursor already marked closed. + # ``closing()``, and its finalizer send CloseOperation and CloseSession. + # Both skip a cursor already marked closed. Wrappers such as SQL Lab's + # row-limited cursor forward the attribute to the driver's cursor. cursor._closed = True # pylint: disable=protected-access + # Metadata reads after a Stop must not reach the abandoned operation, + # even if its transport stack could not be recognised or dropped. + cursor._last_operation_active = False # pylint: disable=protected-access + cursor._last_operation = None # pylint: disable=protected-access + cursor._description = None # pylint: disable=protected-access + # Discard any buffered reply before releasing the connection locally. + try: + cls._reset_transport(cursor, reconnect=False) + except Exception: # pylint: disable=broad-except + # The closed flag is set before this best-effort, local cleanup. + logger.warning("Query %s: dropping the transport failed", query_id) - @staticmethod - def _reset_transport(cursor: Any) -> bool: + @classmethod + def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: """ - Reopen the cursor's Thrift transport with every buffer emptied. + Close the cursor's Thrift transport with every buffer emptied. impyla's ``reconnect()`` only closes and reopens the socket or HTTP connection, while the buffering layers above it keep their read and write buffers. Returns whether every layer was recognised and reset. + + With ``reconnect`` the transport is opened again. Without it a socket + is left closed until impyla's next RPC opens it; an HTTP client is + still reopened, which creates a connection object without network I/O, + because closing it a second time raises. """ # pylint: disable=import-outside-toplevel from io import BytesIO + # The same transport object impyla's own ``reconnect()`` reopens. + transport = cursor.session.service.client._iprot.trans + stack = cls._transport_buffers(transport) + if stack is None: + if not reconnect: + transport.close() + return False + buffers, over_socket = stack + + # Closing drops the socket together with any reply bytes still in it. + try: + transport.close() + except TIME_LIMIT_EXCEPTIONS: + raise + except Exception: # pylint: disable=broad-except + # An HTTP client that was already closed raises on a second close. + if transport.isOpen(): + raise + for buffered_layer, attribute in buffers: + setattr(buffered_layer, attribute, BytesIO()) + if reconnect or not over_socket: + transport.open() + return True + + @staticmethod + def _transport_buffers( + transport: Any, + ) -> tuple[list[tuple[Any, str]], bool] | None: + """ + List the buffers of every layer of an impyla transport stack. + + Returns the (layer, buffer attribute) pairs and whether the stack ends + in a socket rather than an HTTP client, or None if a layer is not + recognised. + """ + # pylint: disable=import-outside-toplevel from impala._thrift_api import ImpalaHttpClient from thrift.transport.TSocket import TSocket from thrift.transport.TTransport import TBufferedTransport @@ -242,60 +302,69 @@ def _reset_transport(cursor: Any) -> bool: except ImportError: pass - # The same transport object impyla's own ``reconnect()`` reopens. - transport = cursor.session.service.client._iprot.trans buffers: list[tuple[Any, str]] = [] layer = transport while not isinstance(layer, TSocket): layout = next((lay for lay in layouts if isinstance(layer, lay[0])), None) if layout is None: - return False + return None _, attributes, inner = layout if not all(hasattr(layer, attribute) for attribute in attributes): - return False + return None buffers.extend((layer, attribute) for attribute in attributes) if inner is None: - break + return buffers, False layer = getattr(layer, inner, None) + return buffers, True - # Closing drops the socket together with any reply bytes still in it. + @staticmethod + def _is_error_reply(ex: Exception) -> bool: + """Whether ``ex`` reports an error reply that was read to its end.""" + # impyla raises HiveServer2Error for an error status, and Thrift raises + # TApplicationException for an exception reply, each after reading the + # whole reply. Any other failure, such as a dropped connection or a + # reply that could not be parsed, can leave part of a reply unread. + # pylint: disable=import-outside-toplevel + reply_errors: list[type[Exception]] = [] try: - transport.close() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - # An HTTP client that was already closed raises on a second close. - if transport.isOpen(): - raise - for buffered_layer, attribute in buffers: - setattr(buffered_layer, attribute, BytesIO()) - transport.open() - return True + from impala.error import HiveServer2Error + + reply_errors.append(HiveServer2Error) + except ImportError: + pass + try: + from thrift.Thrift import TApplicationException + + reply_errors.append(TApplicationException) + except ImportError: + pass + return isinstance(ex, tuple(reply_errors)) @classmethod def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: """Cancel the live operation and release its handles.""" - try: - cursor.cancel_operation() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - logger.warning("Query %s: cancel_operation() failed", query_id) - # The handles are released even when the cancel RPC failed, so a stopped - # query does not leave an operation open on the coordinator for the rest - # of the connection's life. - try: - cursor.close_operation() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - logger.warning("Query %s: close_operation() failed", query_id) - try: - cursor.close() - except TIME_LIMIT_EXCEPTIONS: - raise - except Exception: # pylint: disable=broad-except - logger.warning("Query %s: close() failed", query_id) + # The handles are released even when the coordinator rejects the + # cancel, so a stopped query does not leave an operation open on the + # coordinator for the rest of the connection's life. A call that fails + # without reading its whole reply leaves the transport unusable, so the + # cursor is abandoned instead of reading the next reply from it. + for name, rpc in ( + ("cancel_operation", cursor.cancel_operation), + ("close_operation", cursor.close_operation), + ("close", cursor.close), + ): + try: + rpc() + except TIME_LIMIT_EXCEPTIONS: + # Also covers Stop and the local fetch deadline: their outer + # time-limit handler must not retry this interrupted cleanup. + cls._abandon_cursor(cursor, query_id) + raise + except Exception as ex: # pylint: disable=broad-except + logger.warning("Query %s: %s() failed", query_id, name) + if not cls._is_error_reply(ex): + cls._abandon_cursor(cursor, query_id) + return @classmethod def handle_cursor(cls, cursor: Any, query: Query) -> None: diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index 1aee2fd441f7..c13610c34afd 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -205,15 +205,17 @@ def __getattr__(self, name: str) -> Any: """Delegate metadata and driver-specific methods to the real cursor.""" return getattr(self._cursor, name) - @property - def arraysize(self) -> int: - """Expose the driver's default fetch batch size.""" - return self._cursor.arraysize - - @arraysize.setter - def arraysize(self, value: int) -> None: - """Preserve engine-specific cursor batch-size configuration.""" - self._cursor.arraysize = value + def __setattr__(self, name: str, value: Any) -> None: + """ + Write driver state, such as a closed flag, to the real cursor. + + A flag left on the wrapper would be invisible to the driver when the + caller closes the real cursor. + """ + if name in {"_cursor", "_remaining"}: + object.__setattr__(self, name, value) + else: + setattr(self._cursor, name, value) def fetchmany(self, size: int | None = None) -> list[Any]: """Read no more than the remaining budget, including across batches.""" diff --git a/tests/unit_tests/db_engine_specs/test_hive.py b/tests/unit_tests/db_engine_specs/test_hive.py index b000e1facc7d..36ea92e3236c 100644 --- a/tests/unit_tests/db_engine_specs/test_hive.py +++ b/tests/unit_tests/db_engine_specs/test_hive.py @@ -431,3 +431,52 @@ def test_handle_cursor_propagates_time_limits(timeout: Exception) -> None: with pytest.raises(type(timeout)): HiveEngineSpec.handle_cursor(cursor, query) cursor.fetch_logs.assert_called_once_with() + # The interrupted reply may be partly unread: the cursor's close() sends no + # CloseOperation, and the connection's CloseSession fails before writing. + assert cursor._operationHandle is None + cursor._connection._transport.close.assert_called_once_with() + + +@pytest.mark.parametrize( + "timeout", + [ + SoftTimeLimitExceeded(), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +@pytest.mark.parametrize("interrupted", ["poll", "fetch"]) +def test_fetch_data_abandons_the_cursor_on_time_limits( + interrupted: str, timeout: Exception +) -> None: + """A time limit while fetching sends no cleanup RPC on the transport.""" + from superset.db_engine_specs.base import BaseEngineSpec + from superset.db_engine_specs.hive import HiveEngineSpec + + class ProgrammingError(Exception): + pass + + # The pyhive driver needs thrift, which is not always installed. + ttypes = Mock() + pyhive = Mock(exc=Mock(ProgrammingError=ProgrammingError)) + cursor: Mock = Mock() + if interrupted == "poll": + cursor.poll.side_effect = timeout + with ( + patch.dict( + "sys.modules", + {"pyhive": pyhive, "TCLIService": Mock(ttypes=ttypes)}, + ), + patch.object(BaseEngineSpec, "fetch_data", side_effect=timeout), + ): + with pytest.raises(type(timeout)) as raised: + HiveEngineSpec.fetch_data(cursor) + + assert raised.value is timeout + assert cursor._operationHandle is None + cursor._connection._transport.close.assert_called_once_with() + cursor.close.assert_not_called() diff --git a/tests/unit_tests/db_engine_specs/test_hive_transport.py b/tests/unit_tests/db_engine_specs/test_hive_transport.py new file mode 100644 index 000000000000..b12d6f61f466 --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_hive_transport.py @@ -0,0 +1,203 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Time limits over real pyhive Thrift transports. + +The HiveServer2 stand-in from ``test_impala_transport`` serves pyhive's NOSASL +and SASL PLAIN transports. The time limit is raised right after a reply header +is read, as a signal landing at that point would, so the rest of the reply is +left in the client's buffers. +""" + +from __future__ import annotations + +import socket +import threading +from collections.abc import Iterator +from contextlib import closing +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from celery.exceptions import SoftTimeLimitExceeded + +from superset.db_engine_specs.hive import HiveEngineSpec +from superset.errors import ErrorLevel, SupersetErrorType +from superset.exceptions import SupersetTimeoutException + +pytest.importorskip("pyhive.hive") +pytest.importorskip("thrift_sasl") +transport_tests = pytest.importorskip( + "tests.unit_tests.db_engine_specs.test_impala_transport" +) +TTransportModule = pytest.importorskip("thrift.transport.TTransport") + + +@pytest.fixture(params=["NOSASL", "NONE"], ids=["binary", "sasl"]) +def hive_connection( + request: pytest.FixtureRequest, +) -> Iterator[tuple[Any, Any]]: + """A pyhive connection to a local HiveServer2 over each transport.""" + from pyhive import hive + + handler = transport_tests.FakeHiveServer2() + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.bind(("127.0.0.1", 0)) + listener.listen() + thread = threading.Thread( + target=transport_tests._serve_socket, + args=(handler, listener, request.param == "NONE"), + daemon=True, + ) + thread.start() + conn = hive.connect( + "127.0.0.1", listener.getsockname()[1], auth=request.param, username="user" + ) + # pyhive sets no socket timeout, so a stale read would block forever. + transport = conn._transport + sock = getattr(transport, "_trans", None) or transport._TBufferedTransport__trans + sock.setTimeout(2000) + try: + yield handler, conn + finally: + conn._transport.close() + listener.shutdown(socket.SHUT_RDWR) + listener.close() + thread.join(timeout=2) + + +TIMEOUTS = [ + pytest.param(SoftTimeLimitExceeded(), id="async"), + pytest.param( + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + id="sync", + ), +] + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "entry, interrupted", + [ + ("handle_cursor", "GetOperationStatus"), + ("handle_cursor", "GetLog"), + ("fetch_data", "GetOperationStatus"), + ("fetch_data", "FetchResults"), + ], + ids=["poll", "get_log", "fetch_poll", "rows"], +) +def test_cleanup_sends_nothing_after_an_interrupted_reply( + entry: str, + interrupted: str, + timeout: Exception, + hive_connection: tuple[Any, Any], + app_context: None, +) -> None: + """ + Closing the cursor and the connection must not read leftover reply bytes. + + The executor closes the cursor through ``closing()``, which sends + CloseOperation, and the connection pool closes the connection, which sends + CloseSession. + """ + handler, conn = hive_connection + if entry == "fetch_data": + handler.state = transport_tests.ttypes.TOperationState.FINISHED_STATE + cursor = conn.cursor() + cursor.execute("SELECT 1", async_=True) + transport_tests._interrupt_replies(conn._client._iprot, [(interrupted, timeout)]) + query = Mock(id=1, status="running", progress=0) + + def run() -> None: + if entry == "fetch_data": + HiveEngineSpec.fetch_data(cursor) + else: + HiveEngineSpec.handle_cursor(cursor, query) + + with patch("superset.db_engine_specs.hive.db"): + with pytest.raises(type(timeout)) as raised, closing(cursor): + run() + # Closing a discarded connection is local-only, including in the pool. + conn.close() + + assert raised.value is timeout + assert handler.requests[-1] == interrupted, handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "drop_failure", [OSError("close failed"), SoftTimeLimitExceeded("close")] +) +def test_failed_transport_drop_preserves_timeout_and_disables_session_close( + timeout: Exception, + drop_failure: Exception, + hive_connection: tuple[Any, Any], + app_context: None, +) -> None: + """A failed local close must not restore CloseSession or mask the timeout.""" + handler, conn = hive_connection + cursor = conn.cursor() + cursor.execute("SELECT 1", async_=True) + transport_tests._interrupt_replies( + conn._client._iprot, [("GetOperationStatus", timeout)] + ) + close = conn._transport.close + attempts = 0 + + def interrupted_close() -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise drop_failure + close() + + with patch.object(conn._transport, "close", side_effect=interrupted_close): + with pytest.raises(type(timeout)) as raised, closing(cursor): + HiveEngineSpec.fetch_data(cursor) + conn.close() + + assert raised.value is timeout + assert attempts == 2 + assert handler.requests[-1] == "GetOperationStatus", handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_connection_close_sends_no_rpc_after_an_interrupted_reply( + timeout: Exception, + hive_connection: tuple[Any, Any], + app_context: None, +) -> None: + """Connection.close must not send CloseSession on an interrupted transport.""" + handler, conn = hive_connection + cursor = conn.cursor() + cursor.execute("SELECT 1", async_=True) + transport_tests._interrupt_replies( + conn._client._iprot, [("GetOperationStatus", timeout)] + ) + with pytest.raises(type(timeout)) as raised: + HiveEngineSpec.fetch_data(cursor) + # Isolate connection cleanup from cursor cleanup, even before the fix. + cursor._operationHandle = None + cursor.close() + conn.close() + + assert raised.value is timeout + assert handler.requests[-1] == "GetOperationStatus", handler.requests diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 5e7c469bc921..d417a7a01aec 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -29,6 +29,7 @@ from superset.exceptions import SupersetTimeoutException from superset.models.core import Database from superset.models.sql_lab import Query +from superset.sql.execution.executor import _LimitedCursor from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm from tests.unit_tests.fixtures.common import dttm # noqa: F401 @@ -186,6 +187,10 @@ class _DriverError(Exception): """Represent impyla's optional DB-API error base in unit tests.""" +class _RejectedReplyError(_DriverError): + """Represent impyla's error for an error status, read in full.""" + + @pytest.fixture(autouse=True) def clean_transport() -> Iterator[Mock]: """ @@ -201,7 +206,14 @@ def clean_transport() -> Iterator[Mock]: @pytest.fixture def impala_errors(mocker: MockerFixture) -> None: """Provide driver errors without requiring the optional impyla package.""" - mocker.patch.dict("sys.modules", {"impala.error": Mock(Error=_DriverError)}) + mocker.patch.dict( + "sys.modules", + { + "impala.error": Mock( + Error=_DriverError, HiveServer2Error=_RejectedReplyError + ) + }, + ) class _AsyncCursor: @@ -274,13 +286,12 @@ def test_fetch_data_timeout( def test_fetch_data_deadline_releases_operation_when_cancel_fails( monotonic_mock: Mock, app_context: None, impala_errors: None ) -> None: - """A failed cancel RPC must not mask the deadline or retain operation handles.""" + """A rejected cancel must not mask the deadline or retain operation handles.""" from flask import current_app - from impala.error import Error cursor = Mock() cursor.is_executing.return_value = True - cursor.cancel_operation.side_effect = Error("cancel RPC failed") + cursor.cancel_operation.side_effect = _RejectedReplyError("cancel rejected") with patch.dict(current_app.config, SQLLAB_TIMEOUT=30): with pytest.raises(TimeoutError, match="Timed out waiting"): @@ -293,6 +304,31 @@ def test_fetch_data_deadline_releases_operation_when_cancel_fails( assert monotonic_mock.call_count == 2 +@patch("superset.db_engine_specs.impala.time.monotonic", side_effect=[0, 31]) +def test_fetch_data_deadline_abandons_the_cursor_when_the_cancel_reply_is_unread( + monotonic_mock: Mock, + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """A cancel that failed partway through its reply ends the cleanup.""" + from flask import current_app + + cursor = Mock() + cursor.is_executing.return_value = True + cursor.cancel_operation.side_effect = _DriverError("connection lost") + + with patch.dict(current_app.config, SQLLAB_TIMEOUT=30): + with pytest.raises(TimeoutError, match="Timed out waiting"): + spec.fetch_data(cursor) + + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() + assert cursor._closed is True + # The transport is dropped, not reconnected, so the next RPC opens it. + clean_transport.assert_called_once_with(cursor, reconnect=False) + + @pytest.mark.parametrize("error", [TypeError("bug"), AttributeError("bug")]) @patch.object(spec, "get_dbapi_mapped_exception") def test_fetch_data_preserves_programming_errors( @@ -381,7 +417,7 @@ def test_fetch_data_propagates_time_limit_during_cancel( app_context: None, impala_errors: None, ) -> None: - """A second time limit in cancellation must replace the original time limit.""" + """A second time limit in cancellation must preserve the original time limit.""" original_timeout = SoftTimeLimitExceeded("waiting") cancel_timeout: Exception if cancel_timeout_type is SupersetTimeoutException: @@ -397,10 +433,10 @@ def test_fetch_data_propagates_time_limit_during_cancel( cursor.is_executing.return_value = True cursor.cancel_operation.side_effect = cancel_timeout - with pytest.raises(cancel_timeout_type) as raised: + with pytest.raises(SoftTimeLimitExceeded) as raised: spec.fetch_data(cursor) - assert raised.value is cancel_timeout + assert raised.value is original_timeout cursor.cancel_operation.assert_called_once_with() cursor.fetchall.assert_not_called() @@ -459,7 +495,7 @@ def test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch( calls.attach_mock(clean_transport, "reset_transport") calls.attach_mock(cursor.cancel_operation, "cancel_operation") if cancel_fails: - cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") + cursor.cancel_operation.side_effect = _RejectedReplyError("cancel rejected") timeout = SoftTimeLimitExceeded() with patch.object(BaseEngineSpec, "fetch_data", side_effect=timeout): @@ -495,13 +531,14 @@ def test_handle_cursor_resets_transport_before_cancelling_an_interrupted_poll( timeout: Exception, clean_transport: Mock, app_context: None, + impala_errors: None, ) -> None: """Discard interrupted status/log replies before cancelling and releasing.""" cursor = Mock() cursor.status.return_value = "RUNNING_STATE" getattr(cursor, interrupted_rpc).side_effect = timeout if cancel_fails: - cursor.cancel_operation.side_effect = RuntimeError("cancel RPC failed") + cursor.cancel_operation.side_effect = _RejectedReplyError("cancel rejected") query = Mock(id=1, extra={}, status="running", progress=0) calls = MagicMock() calls.attach_mock(clean_transport, "reset_transport") @@ -571,7 +608,8 @@ def test_time_limit_skips_cancel_rpcs_without_a_clean_transport( spec.handle_cursor(cursor, query) assert raised.value is timeout - clean_transport.assert_called_once_with(cursor) + # The reset is tried once and the transport then dropped. + assert clean_transport.mock_calls == [call(cursor), call(cursor, reconnect=False)] cursor.cancel_operation.assert_not_called() cursor.close_operation.assert_not_called() cursor.close.assert_not_called() @@ -589,7 +627,7 @@ def test_time_limit_during_transport_reset_propagates( app_context: None, impala_errors: None, ) -> None: - """A second time limit while resetting propagates and sends no cancel RPC.""" + """A second limit while resetting must not hide the original time limit.""" from superset.db_engine_specs.base import BaseEngineSpec second_timeout: Exception @@ -606,12 +644,93 @@ def test_time_limit_during_transport_reset_propagates( cursor.is_executing.return_value = False cursor.execution_failed.return_value = False - with patch.object( - BaseEngineSpec, "fetch_data", side_effect=SoftTimeLimitExceeded("fetching") - ): - with pytest.raises(type(second_timeout)) as raised: + original_timeout = SoftTimeLimitExceeded("fetching") + with patch.object(BaseEngineSpec, "fetch_data", side_effect=original_timeout): + with pytest.raises(SoftTimeLimitExceeded) as raised: spec.fetch_data(cursor) - assert raised.value is second_timeout + assert raised.value is original_timeout cursor.cancel_operation.assert_not_called() assert cursor._closed is True + + +@pytest.mark.parametrize( + "cleanup_rpc", ["cancel_operation", "close_operation", "close"] +) +@pytest.mark.parametrize( + "second_timeout", + [ + SoftTimeLimitExceeded("cleanup"), + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out during cleanup", + level=ErrorLevel.ERROR, + ), + ], + ids=["async", "sync"], +) +def test_time_limit_during_cleanup_abandons_the_cursor( + cleanup_rpc: str, + second_timeout: Exception, + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """A second limit inside a cleanup reply leaves the cursor abandoned.""" + from superset.db_engine_specs.base import BaseEngineSpec + + cursor = Mock() + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + getattr(cursor, cleanup_rpc).side_effect = second_timeout + + original_timeout = SoftTimeLimitExceeded("fetching") + with patch.object(BaseEngineSpec, "fetch_data", side_effect=original_timeout): + with pytest.raises(SoftTimeLimitExceeded) as raised: + spec.fetch_data(cursor) + + assert raised.value is original_timeout + assert cursor._closed is True + assert clean_transport.mock_calls == [call(cursor), call(cursor, reconnect=False)] + + +@pytest.mark.parametrize( + "reset_outcome", + [{"return_value": False}, {"side_effect": OSError("connection refused")}], + ids=["unknown_transport", "reopen_fails"], +) +def test_time_limit_abandons_the_cursor_behind_a_row_limit( + reset_outcome: dict[str, Any], + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """The closed flag reaches the cursor that SQL Lab closes, not the wrapper.""" + from superset.db_engine_specs.base import BaseEngineSpec + + clean_transport.configure_mock(**reset_outcome) + cursor = Mock() + cursor.is_executing.return_value = False + cursor.execution_failed.return_value = False + limited = _LimitedCursor(cursor, 10) + + with patch.object( + BaseEngineSpec, "fetch_data", side_effect=SoftTimeLimitExceeded() + ): + with pytest.raises(SoftTimeLimitExceeded): + spec.fetch_data(limited) + + assert cursor._closed is True + assert "_closed" not in vars(limited) + + +def test_closed_cursor_is_not_cancelled_again( + clean_transport: Mock, app_context: None +) -> None: + """An abandoned or closed cursor sends nothing after a later time limit.""" + cursor = Mock(_closed=True) + + spec._cancel_operation_after_time_limit(cursor, 1) + + clean_transport.assert_not_called() + cursor.cancel_operation.assert_not_called() diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 7110d92d012b..6af03012eb0e 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_pending.py +++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py @@ -42,6 +42,19 @@ from superset.utils.core import QueryStatus +class _RejectedReplyError(Exception): + """Represent impyla's error for an error status, read in full.""" + + +@pytest.fixture +def rejected_reply(mocker: MockerFixture) -> type[Exception]: + """Provide impyla's error-status class without requiring impyla.""" + mocker.patch.dict( + "sys.modules", {"impala.error": Mock(HiveServer2Error=_RejectedReplyError)} + ) + return _RejectedReplyError + + @pytest.fixture(autouse=True) def clean_transport() -> Iterator[Mock]: """ @@ -244,12 +257,14 @@ def test_non_progress_logs_keep_polling(log: str) -> None: db.session.commit.assert_called_once_with() -def test_failed_cancel_still_releases_the_operation() -> None: - """A cancel RPC that errors must not leave the operation and cursor open.""" +def test_failed_cancel_still_releases_the_operation( + rejected_reply: type[Exception], +) -> None: + """A rejected cancel must not leave the operation and cursor open.""" query: Mock = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0) cursor: Mock = Mock() cursor.status.return_value = "PENDING_STATE" - cursor.cancel_operation.side_effect = RuntimeError("rpc failed") + cursor.cancel_operation.side_effect = rejected_reply("rpc rejected") with patch("superset.db_engine_specs.impala.db"): ImpalaEngineSpec.handle_cursor(cursor, query) cursor.close_operation.assert_called_once_with() @@ -384,15 +399,18 @@ def test_stop_cleanup_propagates_soft_time_limit( with patch("superset.db_engine_specs.impala.db"): with pytest.raises(SoftTimeLimitExceeded): ImpalaEngineSpec.handle_cursor(cursor, query) - if not repeat_timeout: - cursor.close.assert_called() + # The interrupted cleanup is not retried by the outer time-limit handler. + assert getattr(cursor, method).call_count == 1 + assert cursor._closed is True @pytest.mark.parametrize("method", ["cancel_operation", "close_operation", "close"]) -def test_cancel_operation_tolerates_rpc_failures(method: str) -> None: - """Ordinary RPC failures must not prevent the remaining cleanup calls.""" +def test_cancel_operation_tolerates_rejected_rpcs( + method: str, rejected_reply: type[Exception] +) -> None: + """A rejected cleanup RPC must not prevent the remaining cleanup calls.""" cursor: Mock = Mock() - getattr(cursor, method).side_effect = RuntimeError("rpc failed") + getattr(cursor, method).side_effect = rejected_reply("rpc rejected") ImpalaEngineSpec._cancel_operation(cursor, 1) @@ -401,6 +419,23 @@ def test_cancel_operation_tolerates_rpc_failures(method: str) -> None: cursor.close.assert_called_once_with() +@pytest.mark.parametrize("method", ["cancel_operation", "close_operation", "close"]) +def test_cancel_operation_abandons_the_cursor_after_an_unread_reply( + method: str, clean_transport: Mock +) -> None: + """A cleanup RPC that may have left its reply unread ends the cleanup.""" + methods = ["cancel_operation", "close_operation", "close"] + cursor: Mock = Mock() + getattr(cursor, method).side_effect = RuntimeError("connection lost") + + ImpalaEngineSpec._cancel_operation(cursor, 1) + + for later in methods[methods.index(method) + 1 :]: + getattr(cursor, later).assert_not_called() + assert cursor._closed is True + clean_transport.assert_called_once_with(cursor, reconnect=False) + + def test_pending_refresh_observes_stop() -> None: """Refreshing the existing query must observe a Stop without a second SELECT.""" app: Flask = Flask(__name__) diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py index 4f08696d0812..b67adc93dd96 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_transport.py +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -38,9 +38,11 @@ import pytest from celery.exceptions import SoftTimeLimitExceeded +from superset.constants import QUERY_EARLY_CANCEL_KEY from superset.db_engine_specs.impala import ImpalaEngineSpec as spec # noqa: N813 from superset.errors import ErrorLevel, SupersetErrorType from superset.exceptions import SupersetTimeoutException +from superset.sql.execution.executor import _LimitedCursor hs2 = pytest.importorskip("impala.hiveserver2") ttypes = pytest.importorskip("impala._thrift_gen.TCLIService.ttypes") @@ -48,6 +50,7 @@ TSocketModule = pytest.importorskip("thrift.transport.TSocket") TTransportModule = pytest.importorskip("thrift.transport.TTransport") TBinaryProtocol = pytest.importorskip("thrift.protocol.TBinaryProtocol") +TProtocol = pytest.importorskip("thrift.protocol.TProtocol") THttpServer = pytest.importorskip("thrift.server.THttpServer") OK = ttypes.TStatus(statusCode=ttypes.TStatusCode.SUCCESS_STATUS) @@ -59,6 +62,8 @@ class FakeHiveServer2: """The HiveServer2 RPCs impyla issues for one SQL Lab query.""" state: int = ttypes.TOperationState.RUNNING_STATE + has_more_rows: bool = False + cancel_status: Any = field(default_factory=lambda: OK) calls: list[tuple[str, int, Any]] = field(default_factory=list) # Every request received, including methods this stand-in does not serve. requests: list[str] = field(default_factory=list) @@ -127,7 +132,7 @@ def FetchResults(self, req: Any) -> Any: # noqa: N802 values = [b"row value " * 5] * 500 return ttypes.TFetchResultsResp( status=OK, - hasMoreRows=False, + hasMoreRows=self.has_more_rows, results=ttypes.TRowSet( startRowOffset=0, rows=[], @@ -143,7 +148,7 @@ def FetchResults(self, req: Any) -> Any: # noqa: N802 def CancelOperation(self, req: Any) -> Any: # noqa: N802 self._record("CancelOperation", req) - return ttypes.TCancelOperationResp(status=OK) + return ttypes.TCancelOperationResp(status=self.cancel_status) def CloseOperation(self, req: Any) -> Any: # noqa: N802 self._record("CloseOperation", req) @@ -276,25 +281,36 @@ def hs2_connection( try: yield mode, handler, conn finally: + conn.close() + listener.shutdown(socket.SHUT_RDWR) listener.close() + thread.join(timeout=2) -def _interrupt_after_header(cursor: Any, reply: str, timeout: Exception) -> None: - """Raise ``timeout`` once the header of the first ``reply`` has been read.""" - protocol = cursor.session.service.client._iprot +def _interrupt_replies(protocol: Any, interrupts: list[tuple[str, Exception]]) -> None: + """ + Raise each exception once the header of the next matching reply is read. + + The interrupts fire in order, each on the first matching reply after the + previous one fired. + """ read_message_begin = protocol.readMessageBegin - armed = [True] + pending = list(interrupts) def interrupted_read_message_begin() -> Any: header = read_message_begin() - if armed[0] and header[0] == reply: - armed[0] = False - raise timeout + if pending and header[0] == pending[0][0]: + raise pending.pop(0)[1] return header protocol.readMessageBegin = interrupted_read_message_begin +def _interrupt_after_header(cursor: Any, reply: str, timeout: Exception) -> None: + """Raise ``timeout`` once the header of the first ``reply`` has been read.""" + _interrupt_replies(cursor.session.service.client._iprot, [(reply, timeout)]) + + TIMEOUTS = [ pytest.param(SoftTimeLimitExceeded(), id="async"), pytest.param( @@ -445,6 +461,199 @@ def test_closing_the_cursor_sends_nothing_after_a_skipped_cancel( assert handler.requests[interrupted + 1 :] == [] +RESET_FAILURES = [ + pytest.param({"return_value": False}, id="unknown_transport"), + pytest.param({"side_effect": OSError("transport close failed")}, id="reset_raises"), +] + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize("reset_failure", RESET_FAILURES) +@pytest.mark.parametrize("read", ["rows", "truncation_probe"]) +def test_row_limited_fetch_sends_nothing_after_a_skipped_cancel( + read: str, + timeout: Exception, + reset_failure: dict[str, Any], + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """ + SQL Lab's row-limit wrapper must mark the driver's cursor abandoned. + + The executor closes the real cursor, not the wrapper, through + ``closing()``; a closed flag left on the wrapper would let impyla send + CloseOperation and CloseSession on the transport holding the interrupted + reply. + """ + _, handler, conn = hs2_connection + handler.state = ttypes.TOperationState.FINISHED_STATE + handler.has_more_rows = True + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + limited = _LimitedCursor(cursor, 500) + + with patch.object(spec, "_reset_transport", **reset_failure): + with closing(cursor): + if read == "rows": + _interrupt_after_header(cursor, "FetchResults", timeout) + with pytest.raises(type(timeout)) as raised: + spec.fetch_data(limited) + assert raised.value is timeout + else: + # The first batch fills the budget, so the probe fetches again. + assert len(spec.fetch_data(limited)) == 500 + _interrupt_after_header(cursor, "FetchResults", timeout) + # The probe reports a failed read as a possibly partial result. + assert limited.check_truncated(spec) is True + cursor.__del__() + + assert cursor._closed is True + assert cursor.description is None + with pytest.raises(hs2.ProgrammingError, match="Operation state is not available"): + cursor.get_log() + interrupted = ( + len(handler.requests) - 1 - handler.requests[::-1].index("FetchResults") + ) + assert handler.requests[interrupted + 1 :] == [] + + +SECOND_TIMEOUTS = [ + pytest.param(SoftTimeLimitExceeded("cleanup"), id="async"), + pytest.param( + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Timed out during cleanup", + level=ErrorLevel.ERROR, + ), + id="sync", + ), +] + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize("second_timeout", SECOND_TIMEOUTS) +@pytest.mark.parametrize( + "cleanup_reply", ["CancelOperation", "CloseOperation", "CloseSession"] +) +@pytest.mark.parametrize( + "entry, interrupted", + [ + ("handle_cursor", "GetOperationStatus"), + ("handle_cursor", "GetLog"), + ("fetch_data", "FetchResults"), + ], +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_time_limit_during_cleanup_abandons_the_cursor( + sleep_mock: Mock, + entry: str, + interrupted: str, + timeout: Exception, + cleanup_reply: str, + second_timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """ + A second limit inside a cleanup reply leaves that reply unread too. + + Nothing may be sent after it: not by the cancellation helper, the + caller's ``closing()``, nor the cursor's finalizer. + """ + _, handler, conn = hs2_connection + if entry == "fetch_data": + handler.state = ttypes.TOperationState.FINISHED_STATE + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + _interrupt_replies( + cursor.session.service.client._iprot, + [ + (interrupted, timeout), + (cleanup_reply, second_timeout), + ], + ) + query = Mock(id=1, extra={}, status="running", progress=0) + + def run() -> None: + if entry == "fetch_data": + spec.fetch_data(cursor) + else: + spec.handle_cursor(cursor, query) + + with patch("superset.db_engine_specs.impala.db"): + with pytest.raises(type(timeout)) as raised, closing(cursor): + run() + cursor.__del__() + + assert raised.value is timeout + assert cursor._closed is True + assert handler.requests[-1] == cleanup_reply, handler.requests + + +@pytest.mark.parametrize( + "cleanup_reply", ["CancelOperation", "CloseOperation", "CloseSession"] +) +@patch("superset.db_engine_specs.impala.time.sleep") +def test_failed_stop_cleanup_rpc_abandons_the_cursor( + sleep_mock: Mock, + cleanup_reply: str, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """ + A cleanup RPC that fails partway through its reply ends the cleanup. + + Its unread rest must not be read as the next reply. Metadata reads after + a Stop and further polls must not send RPCs for the abandoned operation. + """ + _, handler, conn = hs2_connection + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + # A reply that cannot be parsed, as leftover bytes read as a header would + # be, fails without impyla retrying the call. + _interrupt_after_header( + cursor, cleanup_reply, TProtocol.TProtocolException(message="bad reply") + ) + query = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0) + + with patch("superset.db_engine_specs.impala.db"): + with closing(cursor): + spec.handle_cursor(cursor, query) + cursor.__del__() + + assert cursor._closed is True + assert handler.requests[-1] == cleanup_reply, handler.requests + assert cursor.description is None + with pytest.raises(hs2.ProgrammingError, match="Operation state is not available"): + cursor.status() + assert handler.requests[-1] == cleanup_reply, handler.requests + + +@patch("superset.db_engine_specs.impala.time.sleep") +def test_rejected_cancel_still_releases_the_operation( + sleep_mock: Mock, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, + caplog: pytest.LogCaptureFixture, +) -> None: + """An error status is a whole reply, so the handles are still released.""" + _, handler, conn = hs2_connection + handler.cancel_status = ttypes.TStatus( + statusCode=ttypes.TStatusCode.ERROR_STATUS, errorMessage="already done" + ) + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + query = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0) + + with patch("superset.db_engine_specs.impala.db"): + spec.handle_cursor(cursor, query) + + methods = [method for method, _, _ in handler.calls] + assert methods[-3:] == ["CancelOperation", "CloseOperation", "CloseSession"] + assert cursor._closed is True + assert "cancel_operation() failed" in caplog.text + + def test_reset_transport_refuses_an_unknown_transport() -> None: """An unrecognised buffering layer is left untouched and reported unclean.""" cursor = Mock() diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py index 0ecd0109149c..df3f9062f0d5 100644 --- a/tests/unit_tests/sql/execution/test_executor.py +++ b/tests/unit_tests/sql/execution/test_executor.py @@ -1965,6 +1965,28 @@ def test_execute_bounds_cursor_fetches( ] +def test_limited_cursor_writes_driver_state_to_the_real_cursor() -> None: + """ + State an engine spec sets, such as a closed flag, reaches the real cursor. + + The executor closes the real cursor, so a flag left on the wrapper would + be invisible to the driver. + """ + from superset.sql.execution.executor import _LimitedCursor + + cursor = create_mock_cursor(["n"], [(n,) for n in range(3)]) + cursor.fetchmany.return_value = [(0,)] + limited = _LimitedCursor(cursor, 3) + probe = _LimitedCursor(limited, 1) + + probe._closed = True + + assert cursor._closed is True + assert "_closed" not in vars(limited) + assert "_closed" not in vars(probe) + assert probe.fetchall() == [(0,)] + + @pytest.mark.parametrize("first_read", ["fetchmany", "fetchone", "iterate"]) def test_limited_cursor_shares_read_budget(first_read: str) -> None: """All cursor read methods share a budget; metadata stays on the driver.""" From 9c6a37308ca521a4ffc665202250d0626cbdbef4 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 17:33:21 +0000 Subject: [PATCH 11/14] fix(sqllab): abandon interrupted HiveServer2 calls at the driver boundary --- .../configuration/configuring-superset.mdx | 26 +- superset/db_engine_specs/base.py | 41 +++- superset/db_engine_specs/hive.py | 106 +++++---- superset/db_engine_specs/impala.py | 146 +++++++----- superset/models/core.py | 9 +- superset/sql/execution/executor.py | 5 + superset/sql_lab.py | 11 + superset/utils/hive.py | 125 ++++++++++ .../db_engine_specs/test_hive_transport.py | 19 ++ .../unit_tests/db_engine_specs/test_impala.py | 46 ++-- .../db_engine_specs/test_impala_transport.py | 225 +++++++++++++++--- 11 files changed, 570 insertions(+), 189 deletions(-) create mode 100644 superset/utils/hive.py diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 6e164309165a..c15be805d124 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -833,21 +833,17 @@ 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. 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. +RPC calls. A signal-based SQL Lab time limit can interrupt any HiveServer2 +response, including execution, lazy result metadata, status, logs, rows, or +cleanup. Superset abandons the affected Impala or Hive cursor and connection +before driver exception handlers run: subsequent cancellation, close, and +finalizer calls send no RPCs on that transport. The socket is dropped locally; +cleanup failures are logged without replacing the original time-limit exception. +No reconnect or cancellation is attempted after a signal-based limit. Configure +server-side query and session timeouts to release an abandoned operation. +The local Impala polling deadline and an explicit Stop still request cancellation +when the transport has not been interrupted; an incomplete cleanup reply also +abandons the transport rather than attempting another cleanup RPC. ## SQL Lab Query History Pruning diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 1d5a93436591..7faedcf543f2 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -21,6 +21,8 @@ import logging import re import warnings +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime, timedelta from inspect import signature from re import Match, Pattern @@ -141,7 +143,13 @@ # 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) +# TimerTimeout raises KeyboardInterrupt inside its block on Windows and converts +# it to SupersetTimeoutException on exit, after inner cursor cleanup has run. +TIME_LIMIT_EXCEPTIONS = ( + SoftTimeLimitExceeded, + SupersetTimeoutException, + KeyboardInterrupt, +) def convert_inspector_columns(cols: list[SQLAColumnType]) -> list[ResultSetColumnType]: @@ -2071,6 +2079,37 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: for handling the cursor and updating progress information in the query object""" + @classmethod + def prepare_connection(cls, connection: Any) -> None: + """Install driver-specific safeguards before creating a cursor.""" + + @classmethod + def prepare_cursor(cls, cursor: Any) -> Callable[[BaseException], None] | None: + """Install safeguards and return the driver's timeout-abandon callback.""" + return None + + @classmethod + def _abandon_cursor(cls, cursor: Any) -> None: + """Disable driver cleanup when an interrupted reply cannot be reused.""" + + @classmethod + @contextmanager + def cursor_context(cls, cursor: Any) -> Iterator[None]: + """Preserve time limits before callers or drivers start cursor cleanup.""" + on_timeout = None + try: + on_timeout = cls.prepare_cursor(cursor) + yield + except TIME_LIMIT_EXCEPTIONS as ex: + try: + if on_timeout is not None: + on_timeout(ex) + else: + cls._abandon_cursor(cursor) + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except + logger.warning("Abandoning an interrupted cursor failed") + raise + @classmethod # pylint: disable=consider-using-transaction def execute_with_cursor( diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 40bb7bcd64d3..261f6e844896 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -21,6 +21,7 @@ import re import tempfile import time +from collections.abc import Callable from datetime import datetime from typing import Any, TYPE_CHECKING from urllib import parse @@ -50,6 +51,12 @@ from superset.models.sql_lab import Query from superset.sql.parse import Table from superset.superset_typing import ResultSetColumnType +from superset.utils.hive import ( + guard_cursor, + guard_cursor_factory, + real_cursor, + ThriftClientGuard, +) if TYPE_CHECKING: from superset.models.core import Database @@ -197,44 +204,61 @@ def patch(cls) -> None: hive.ttypes = patched_ttypes @classmethod + @guard_cursor def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]: # pylint: disable=import-outside-toplevel 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: - 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 + return super().fetch_data(cursor, limit) + except pyhive.exc.ProgrammingError: + return [] + + @classmethod + def prepare_connection(cls, connection: Any) -> None: + """Guard every RPC on the PyHive connection's shared Thrift client.""" + if not isinstance(connection._client, ThriftClientGuard): # pylint: disable=protected-access + connection._client = ThriftClientGuard(connection._client) # pylint: disable=protected-access + connection._client.register(connection, cls._abandon_connection) # pylint: disable=protected-access + guard_cursor_factory(connection, cls.prepare_cursor) + + @classmethod + def prepare_cursor(cls, cursor: Any) -> Callable[[BaseException], None] | None: + """Register real cursor cleanup before executing any driver call.""" + cursor = real_cursor(cursor) + connection = cursor._connection # pylint: disable=protected-access + cls.prepare_connection(connection) + connection._client.register(cursor, cls._abandon_cursor) # pylint: disable=protected-access + return connection._client.abandon # pylint: disable=protected-access @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. - """ + def _abandon_connection(connection: Any) -> None: + """Replace CloseSession with a best-effort local transport drop.""" + if getattr(connection, "_superset_abandoned", False) is True: + return + connection._superset_abandoned = True # pylint: disable=protected-access + + def close() -> None: + try: + connection._transport.close() # pylint: disable=protected-access + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except + logger.warning("Closing the transport after a time limit failed") + + connection.close = close + close() + + @classmethod + def _abandon_cursor(cls, cursor: Any) -> None: + """Disable cursor and connection RPC cleanup before dropping the socket.""" + cursor = real_cursor(cursor) 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") + cls._abandon_connection(cursor._connection) # pylint: disable=protected-access @classmethod def df_to_sql( @@ -457,13 +481,10 @@ def get_tracking_url_from_logs(cls, log_lines: list[str]) -> str | None: return None @classmethod + @guard_cursor 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 + cls._poll_cursor(cursor, query) @classmethod def _poll_cursor( # pylint: disable=too-many-locals # noqa: C901 @@ -497,7 +518,7 @@ def _poll_cursor( # pylint: disable=too-many-locals # noqa: C901 log = "\n".join(logs) if logs else "" except TIME_LIMIT_EXCEPTIONS: raise - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.warning("Call to GetLog() failed") log = "" @@ -567,7 +588,7 @@ def where_latest_partition( table, show_first=True, ) - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except # table is not partitioned return None if values is not None and columns is not None: @@ -668,14 +689,17 @@ def impersonate_user( return url, engine_kwargs - @staticmethod - def execute( # type: ignore - cursor, + @classmethod + @guard_cursor + def execute( + cls, + cursor: Any, query: str, database: Database, async_: bool = False, - ): # pylint: disable=arguments-differ - kwargs = {"async": async_} + **kwargs: Any, + ) -> None: + kwargs["async"] = async_ cursor.execute(query, **kwargs) @classmethod diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 02a5fd8421c5..98a43e31bb25 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -20,6 +20,7 @@ import logging import re import time +from collections.abc import Callable from datetime import datetime from typing import Any, Optional, TYPE_CHECKING @@ -36,6 +37,12 @@ ) from superset.models.sql_lab import Query from superset.utils.core import QueryStatus +from superset.utils.hive import ( + guard_cursor, + guard_cursor_factory, + real_cursor, + ThriftClientGuard, +) from superset.utils.network import get_ssrf_safe_requester, is_safe_host if TYPE_CHECKING: @@ -124,6 +131,7 @@ def has_implicit_cancel(cls) -> bool: return False @classmethod + @guard_cursor def execute( cls, cursor: Any, @@ -133,75 +141,84 @@ def execute( ) -> None: try: cursor.execute_async(query) + except TIME_LIMIT_EXCEPTIONS: + raise except Exception as ex: raise cls.get_dbapi_mapped_exception(ex) from ex @classmethod + @guard_cursor def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, ...]]: """Wait for asynchronous operations using the public cursor API.""" - try: - if callable(getattr(cursor, "execute_async", None)): - from impala.error import Error - - deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"] - try: - while cursor.is_executing(): - if time.monotonic() >= deadline: - cls._cancel_operation(cursor) - raise TimeoutError( - "Timed out waiting for the Impala operation" - ) - time.sleep(0.1) - if cursor.execution_failed(): - # Fetching surfaces the driver's detailed asynchronous error, - # including errors on statements that produce no result set. - cursor.fetchall() - except Error as ex: - raise cls.get_dbapi_mapped_exception(ex) from ex - return super().fetch_data(cursor, limit) - except TIME_LIMIT_EXCEPTIONS: - cls._cancel_operation_after_time_limit(cursor) - raise + if callable(getattr(cursor, "execute_async", None)): + from impala.error import Error + + deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"] + try: + while cursor.is_executing(): + if time.monotonic() >= deadline: + cls._cancel_operation(cursor) + raise TimeoutError("Timed out waiting for the Impala operation") + time.sleep(0.1) + if cursor.execution_failed(): + # Fetching surfaces the driver's detailed asynchronous error, + # including errors on statements that produce no result set. + cursor.fetchall() + except Error as ex: + raise cls.get_dbapi_mapped_exception(ex) from ex + return super().fetch_data(cursor, limit) @classmethod - def _cancel_operation_after_time_limit( - cls, cursor: Any, query_id: int | None = None - ) -> None: - """Cancel and release the operation on a transport with no stale bytes.""" - # A closed cursor has no operation left to cancel, and an abandoned one - # must send nothing more. - if getattr(cursor, "_closed", False) is True: - return - # A signal can interrupt any RPC reply, including status, logs, or rows. - # The unread rest of that reply stays in the socket and in the - # transport's read buffer, which survives a plain close and reopen, so - # the cancel reply would be parsed from it and could hang or request a - # garbage-sized read. Without a transport known to be clean, no RPC is - # sent and Impala's own timeouts close the operation. - try: - reset = cls._reset_transport(cursor) - except Exception: # pylint: disable=broad-except - # Cleanup must not replace the original time-limit exception, - # including when a second signal interrupts the reset itself. - logger.warning("Query %s: resetting the transport failed", query_id) - reset = False - if not reset: - logger.warning( - "Query %s: not cancelling the operation on an unclean transport", - query_id, - ) - cls._abandon_cursor(cursor, query_id) - return - try: - cls._cancel_operation(cursor, query_id) - except TIME_LIMIT_EXCEPTIONS: - # Another limit landed inside a cleanup reply, so the transport - # holds unread bytes again. - logger.warning("Query %s: time limit during cancellation cleanup", query_id) + def prepare_connection(cls, connection: Any) -> None: + """Guard OpenSession and all cursors sharing this connection's client.""" + service = connection.service + if not isinstance(service.client, ThriftClientGuard): + service.client = ThriftClientGuard(service.client) + service.client.register(service, cls._abandon_service) + guard_cursor_factory(connection, cls.prepare_cursor) + + @staticmethod + def _abandon_service(service: Any) -> None: + """Make connection close local-only and non-fatal after interruption.""" + transport = service.client._iprot.trans + + def close() -> None: + try: + transport.close() + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except + logger.warning("Dropping the interrupted Impala transport failed") + + service.close = close + close() + + @classmethod + def prepare_cursor(cls, cursor: Any) -> Callable[[BaseException], None] | None: + """Cover lazy metadata and cleanup inside the driver, as well as rows.""" + cursor = real_cursor(cursor) + session = getattr(cursor, "session", None) + if session is None: + return None + service = session.service + if not isinstance(service.client, ThriftClientGuard): + service.client = ThriftClientGuard(service.client) + guard = service.client + session.client = guard + if cursor._last_operation is not None: # pylint: disable=protected-access + cursor._last_operation.client = guard # pylint: disable=protected-access + guard.register(service, cls._abandon_service) + guard.register(cursor, cls._abandon_cursor) + return guard.abandon @classmethod def _abandon_cursor(cls, cursor: Any, query_id: int | None = None) -> None: """Send no further RPC on a transport that may hold an unread reply.""" + cursor = real_cursor(cursor) + if ( + getattr(cursor, "_superset_abandoned", False) is True + or getattr(cursor, "_closed", False) is True + ): + return + cursor._superset_abandoned = True # pylint: disable=protected-access # ``HiveServer2Cursor.close()``, which callers reach through # ``closing()``, and its finalizer send CloseOperation and CloseSession. # Both skip a cursor already marked closed. Wrappers such as SQL Lab's @@ -215,7 +232,7 @@ def _abandon_cursor(cls, cursor: Any, query_id: int | None = None) -> None: # Discard any buffered reply before releasing the connection locally. try: cls._reset_transport(cursor, reconnect=False) - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except # The closed flag is set before this best-effort, local cleanup. logger.warning("Query %s: dropping the transport failed", query_id) @@ -250,7 +267,7 @@ def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: transport.close() except TIME_LIMIT_EXCEPTIONS: raise - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except # An HTTP client that was already closed raises on a second close. if transport.isOpen(): raise @@ -367,6 +384,7 @@ def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: return @classmethod + @guard_cursor def handle_cursor(cls, cursor: Any, query: Query) -> None: """Stop query and updates progress information""" @@ -408,7 +426,7 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: log = cursor.get_log() or "" except TIME_LIMIT_EXCEPTIONS: raise - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.warning("Call to GetLog() failed") log = "" @@ -424,12 +442,8 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: time.sleep(sleep_interval) status = cursor.status() except TIME_LIMIT_EXCEPTIONS: - # SQL Lab ends the query as timed out once this propagates; the - # operation is cancelled first so it does not keep running on the - # coordinator while fetch_data() waits on it. - cls._cancel_operation_after_time_limit(cursor, query_id) raise - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.debug("Call to status() failed ") return @@ -527,7 +541,7 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: ] ) response = requester.post(url, timeout=3, allow_redirects=False) - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except return False return bool(response and response.status_code == 200) diff --git a/superset/models/core.py b/superset/models/core.py index c6a34ab2f45e..f93c8f0321c0 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -671,6 +671,7 @@ def run_prequeries( dbapi_connection: Any, connection_record: Any, # pylint: disable=unused-argument ) -> None: + self.db_engine_spec.prepare_connection(dbapi_connection) cursor = dbapi_connection.cursor() try: from superset.sql.execution.cancellation import ( @@ -679,8 +680,11 @@ def run_prequeries( query_executed, ) - with cancellable_cursor( - self, cursor, catalog, schema + with ( + self.db_engine_spec.cursor_context(cursor), + cancellable_cursor( + self, cursor, catalog, schema + ), ): for prequery in prequeries: check_query_deadline() @@ -851,6 +855,7 @@ def open_connection() -> Iterator[Connection]: ) as engine: with check_for_oauth2(self): with closing(engine.raw_connection()) as conn: + self.db_engine_spec.prepare_connection(conn.dbapi_connection) yield conn with self._open_with_oauth2_retry(open_connection) as conn: diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index c13610c34afd..b6ace98a0ab0 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -72,6 +72,7 @@ from flask_babel import gettext as __ from superset import db +from superset.db_engine_specs.base import TIME_LIMIT_EXCEPTIONS from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import ( OAuth2Error, @@ -239,6 +240,8 @@ def check_truncated(self, db_engine_spec: type[BaseEngineSpec]) -> bool: return False try: return bool(db_engine_spec.fetch_data(_LimitedCursor(self._cursor, 1))) + except TIME_LIMIT_EXCEPTIONS: + raise except Exception: # pylint: disable=broad-except logger.warning( "Truncation probe failed; reporting the result as partial", @@ -920,6 +923,7 @@ def execute() -> list[tuple[str, SupersetResultSet | None, float, int]]: with ( contextlib.closing(conn.cursor()) as cursor, + self.database.db_engine_spec.cursor_context(cursor), cancellable_cursor(self.database, cursor, catalog, schema), ): return execute_sql_with_cursor( @@ -1646,6 +1650,7 @@ def _cancel_query(database: Database, query: Any) -> bool: source=QuerySource.SQL_LAB, ) as engine: with contextlib.closing(engine.raw_connection()) as conn: + database.db_engine_spec.prepare_connection(conn.dbapi_connection) with contextlib.closing(conn.cursor()) as cursor: return database.db_engine_spec.cancel_query( cursor, query, cancel_query_id diff --git a/superset/sql_lab.py b/superset/sql_lab.py index bd4bcea61407..ae48811e924c 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -330,6 +330,16 @@ def execute_query( # pylint: disable=too-many-statements, too-many-locals # no database: Database = query.database db_engine_spec = database.db_engine_spec + with db_engine_spec.cursor_context(cursor): + return _execute_query(query, cursor, log_params) + + +def _execute_query( + query: Query, cursor: Any, log_params: Optional[dict[str, Any]] = None +) -> SupersetResultSet: + """Run a statement inside the engine's interruption-safe cursor scope.""" + database = query.database + db_engine_spec = database.db_engine_spec try: log_query = app.config["QUERY_LOGGER"] if log_query: @@ -903,6 +913,7 @@ def cancel_query(query: Query) -> bool: source=QuerySource.SQL_LAB, ) as engine: with closing(engine.raw_connection()) as conn: + query.database.db_engine_spec.prepare_connection(conn.dbapi_connection) with closing(conn.cursor()) as cursor: return query.database.db_engine_spec.cancel_query( cursor, query, cancel_query_id diff --git a/superset/utils/hive.py b/superset/utils/hive.py new file mode 100644 index 000000000000..5f7158bc7423 --- /dev/null +++ b/superset/utils/hive.py @@ -0,0 +1,125 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Fail-closed guards for signal-interrupted HiveServer2 driver calls.""" + +from __future__ import annotations + +import logging +import weakref +from collections.abc import Callable +from functools import wraps +from typing import Any, cast, TypeVar + +from superset.db_engine_specs.base import TIME_LIMIT_EXCEPTIONS + +logger = logging.getLogger(__name__) +GuardedMethod = TypeVar("GuardedMethod", bound=Callable[..., Any]) + + +class ThriftClientGuard: + """ + Guard every generated Thrift RPC, not a list of individual RPC methods. + + A signal can leave a reply partially read. Abandon targets before returning + control to driver exception handlers, which can themselves issue cleanup + RPCs. Subsequent RPC attempts re-raise the original exception without I/O. + Only this client's instances are wrapped; driver classes are not patched. + """ + + def __init__(self, client: Any) -> None: + """Keep the original client and weakly tracked cleanup targets.""" + self._client = client + self._targets: dict[ + int, tuple[weakref.ReferenceType[Any], Callable[[Any], None]] + ] = {} + self._exception: BaseException | None = None + + def register(self, target: Any, abandon: Callable[[Any], None]) -> None: + """Track a real cursor or connection without extending its lifetime.""" + key = id(target) + owner = weakref.ref(self) + + def discard(_ref: weakref.ReferenceType[Any]) -> None: + if (guard := owner()) is not None: + guard._targets.pop(key, None) + + self._targets[key] = (weakref.ref(target, discard), abandon) + + def __getattr__(self, name: str) -> Any: + """Forward attributes and wrap all callable driver-client entry points.""" + value = getattr(self._client, name) + if not callable(value): + return value + + @wraps(value) + def call(*args: Any, **kwargs: Any) -> Any: + if self._exception is not None: + raise self._exception + try: + return value(*args, **kwargs) + except TIME_LIMIT_EXCEPTIONS as ex: + self.abandon(ex) + raise + + return call + + def abandon(self, exception: BaseException) -> None: + """Disable all further RPCs before attempting best-effort local cleanup.""" + if self._exception is not None: + return + self._exception = exception + for ref, abandon in list(self._targets.values()): + if (target := ref()) is not None: + try: + abandon(target) + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except + logger.warning("Abandoning an interrupted Thrift target failed") + + +def guard_cursor(method: GuardedMethod) -> GuardedMethod: + """Apply the engine's shared timeout scope to direct spec entry points.""" + + @wraps(method) + def guarded(cls: Any, cursor: Any, *args: Any, **kwargs: Any) -> Any: + with cls.cursor_context(cursor): + return method(cls, cursor, *args, **kwargs) + + return cast(GuardedMethod, guarded) + + +def real_cursor(cursor: Any) -> Any: + """Unwrap cursor adapters so state and weak references belong to the driver.""" + while "_cursor" in vars(cursor): + cursor = vars(cursor)["_cursor"] + return cursor + + +def guard_cursor_factory(connection: Any, prepare: Callable[[Any], object]) -> None: + """Register cursors before exposing them to callers, including direct users.""" + if vars(connection).get("_superset_cursor_guard") is True: + return + cursor_factory = connection.cursor + + @wraps(cursor_factory) + def cursor(*args: Any, **kwargs: Any) -> Any: + result = cursor_factory(*args, **kwargs) + prepare(result) + return result + + connection.cursor = cursor + connection._superset_cursor_guard = True # pylint: disable=protected-access diff --git a/tests/unit_tests/db_engine_specs/test_hive_transport.py b/tests/unit_tests/db_engine_specs/test_hive_transport.py index b12d6f61f466..902392769b40 100644 --- a/tests/unit_tests/db_engine_specs/test_hive_transport.py +++ b/tests/unit_tests/db_engine_specs/test_hive_transport.py @@ -201,3 +201,22 @@ def test_connection_close_sends_no_rpc_after_an_interrupted_reply( assert raised.value is timeout assert handler.requests[-1] == "GetOperationStatus", handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_execute_abandons_before_cursor_and_connection_cleanup( + timeout: Exception, + hive_connection: tuple[Any, Any], + app_context: None, +) -> None: + """An interrupted ExecuteStatement must not be followed by CloseSession.""" + handler, conn = hive_connection + cursor = conn.cursor() + transport_tests._interrupt_replies( + conn._client._iprot, [("ExecuteStatement", timeout)] + ) + with pytest.raises(type(timeout)) as raised, closing(cursor): + HiveEngineSpec.execute(cursor, "SELECT 1", None, async_=True) # type: ignore[arg-type] + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == "ExecuteStatement", handler.requests diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index d417a7a01aec..296431eb0a23 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -194,7 +194,7 @@ class _RejectedReplyError(_DriverError): @pytest.fixture(autouse=True) def clean_transport() -> Iterator[Mock]: """ - Report the transport reset as successful for mock cursors. + Stub local transport drops for mock cursors. Mock cursors have no Thrift transport; resets over real impyla transports are covered in ``test_impala_transport.py``. @@ -394,16 +394,16 @@ def test_engines_without_a_sensible_default_offer_none() -> None: ids=["async", "sync"], ) @patch("superset.db_engine_specs.impala.time.sleep") -def test_fetch_data_cancels_the_operation_on_time_limit( +def test_fetch_data_abandons_the_operation_on_time_limit( sleep_mock: Mock, timeout: Exception, app_context: None, impala_errors: None ) -> None: - """A SQL Lab time limit hitting the wait must not leave the operation running.""" + """A SQL Lab time limit disables RPC cleanup before propagating.""" sleep_mock.side_effect = timeout cursor = Mock() cursor.is_executing.return_value = True with pytest.raises(type(timeout)): spec.fetch_data(cursor) - cursor.cancel_operation.assert_called_once_with() + cursor.cancel_operation.assert_not_called() cursor.fetchall.assert_not_called() @@ -417,7 +417,7 @@ def test_fetch_data_propagates_time_limit_during_cancel( app_context: None, impala_errors: None, ) -> None: - """A second time limit in cancellation must preserve the original time limit.""" + """An expired wait never enters RPC cancellation, even if it would time out.""" original_timeout = SoftTimeLimitExceeded("waiting") cancel_timeout: Exception if cancel_timeout_type is SupersetTimeoutException: @@ -437,7 +437,7 @@ def test_fetch_data_propagates_time_limit_during_cancel( spec.fetch_data(cursor) assert raised.value is original_timeout - cursor.cancel_operation.assert_called_once_with() + cursor.cancel_operation.assert_not_called() cursor.fetchall.assert_not_called() @@ -454,7 +454,7 @@ def test_fetch_data_propagates_time_limit_during_cancel( ], ids=["async", "sync"], ) -def test_fetch_data_cancels_on_time_limit_during_row_fetch( +def test_fetch_data_abandons_on_time_limit_during_row_fetch( timeout: Exception, cancel_fails: bool, app_context: None, @@ -475,17 +475,17 @@ def test_fetch_data_cancels_on_time_limit_during_row_fetch( assert raised.value is timeout fetch.assert_called_once_with(cursor, 10) - cursor.cancel_operation.assert_called_once_with() + cursor.cancel_operation.assert_not_called() @pytest.mark.parametrize("cancel_fails", [False, True]) -def test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch( +def test_fetch_data_drops_transport_without_cancelling_an_interrupted_fetch( cancel_fails: bool, clean_transport: Mock, app_context: None, impala_errors: None, ) -> None: - """Discard an interrupted response before cancellation and release handles.""" + """Abandon an interrupted response without issuing any cleanup RPC.""" from superset.db_engine_specs.base import BaseEngineSpec cursor = Mock() @@ -503,12 +503,9 @@ def test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch( spec.fetch_data(cursor) assert raised.value is timeout - assert calls.mock_calls[:2] == [ - call.reset_transport(cursor), - call.cancel_operation(), - ] - cursor.close_operation.assert_called_once_with() - cursor.close.assert_called_once_with() + assert calls.mock_calls == [call.reset_transport(cursor, reconnect=False)] + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() @pytest.mark.parametrize("interrupted_rpc", ["status", "get_log"]) @@ -525,7 +522,7 @@ def test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch( ], ids=["async", "sync"], ) -def test_handle_cursor_resets_transport_before_cancelling_an_interrupted_poll( +def test_handle_cursor_drops_transport_without_cancelling_an_interrupted_poll( interrupted_rpc: str, cancel_fails: bool, timeout: Exception, @@ -552,10 +549,7 @@ def test_handle_cursor_resets_transport_before_cancelling_an_interrupted_poll( assert raised.value is timeout assert calls.mock_calls == [ - call.reset_transport(cursor), - call.cancel_operation(), - call.close_operation(), - call.close(), + call.reset_transport(cursor, reconnect=False), ] @@ -608,8 +602,8 @@ def test_time_limit_skips_cancel_rpcs_without_a_clean_transport( spec.handle_cursor(cursor, query) assert raised.value is timeout - # The reset is tried once and the transport then dropped. - assert clean_transport.mock_calls == [call(cursor), call(cursor, reconnect=False)] + # The transport is dropped once without reconnecting. + assert clean_transport.mock_calls == [call(cursor, reconnect=False)] cursor.cancel_operation.assert_not_called() cursor.close_operation.assert_not_called() cursor.close.assert_not_called() @@ -676,7 +670,7 @@ def test_time_limit_during_cleanup_abandons_the_cursor( app_context: None, impala_errors: None, ) -> None: - """A second limit inside a cleanup reply leaves the cursor abandoned.""" + """An expired fetch never enters RPC cleanup, even if it would time out.""" from superset.db_engine_specs.base import BaseEngineSpec cursor = Mock() @@ -691,7 +685,7 @@ def test_time_limit_during_cleanup_abandons_the_cursor( assert raised.value is original_timeout assert cursor._closed is True - assert clean_transport.mock_calls == [call(cursor), call(cursor, reconnect=False)] + assert clean_transport.mock_calls == [call(cursor, reconnect=False)] @pytest.mark.parametrize( @@ -730,7 +724,7 @@ def test_closed_cursor_is_not_cancelled_again( """An abandoned or closed cursor sends nothing after a later time limit.""" cursor = Mock(_closed=True) - spec._cancel_operation_after_time_limit(cursor, 1) + spec._abandon_cursor(cursor, 1) clean_transport.assert_not_called() cursor.cancel_operation.assert_not_called() diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py index b67adc93dd96..a597c24e1613 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_transport.py +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -25,7 +25,6 @@ from __future__ import annotations -import logging import socket import struct import threading @@ -287,7 +286,9 @@ def hs2_connection( thread.join(timeout=2) -def _interrupt_replies(protocol: Any, interrupts: list[tuple[str, Exception]]) -> None: +def _interrupt_replies( + protocol: Any, interrupts: list[tuple[str, BaseException]] +) -> None: """ Raise each exception once the header of the next matching reply is read. @@ -306,7 +307,7 @@ def interrupted_read_message_begin() -> Any: protocol.readMessageBegin = interrupted_read_message_begin -def _interrupt_after_header(cursor: Any, reply: str, timeout: Exception) -> None: +def _interrupt_after_header(cursor: Any, reply: str, timeout: BaseException) -> None: """Raise ``timeout`` once the header of the first ``reply`` has been read.""" _interrupt_replies(cursor.session.service.client._iprot, [(reply, timeout)]) @@ -324,36 +325,12 @@ def _interrupt_after_header(cursor: Any, reply: str, timeout: Exception) -> None ] -def _assert_cancelled_cleanly( - mode: str, +def _assert_abandoned( handler: FakeHiveServer2, interrupted: str, - caplog: pytest.LogCaptureFixture, ) -> None: - methods = [method for method, _, _ in handler.calls] - # A cancel that read leftover bytes as its reply blocks until the socket - # timeout, after which impyla sends it again on a reopened transport. - assert methods.count("CancelOperation") == 1, methods - cancel_index = methods.index("CancelOperation") - _, cancel_connection, cancel = handler.calls[cancel_index] - assert cancel.operationHandle.operationId.guid == OPERATION_GUID - assert "CloseOperation" in methods[cancel_index:], methods - interrupted_connection = max( - connection - for method, connection, _ in handler.calls[:cancel_index] - if method == interrupted - ) - if mode != "http": - # The interrupted reply's connection is dropped, not reused. - assert cancel_connection > interrupted_connection - # Every cancellation RPC read its own reply. - failures = [ - record.getMessage() - for record in caplog.records - if record.name == "superset.db_engine_specs.impala" - and record.levelno >= logging.WARNING - ] - assert failures == [] + """The interrupted reply must be the last RPC, without reconnection.""" + assert handler.requests[-1] == interrupted, handler.requests @pytest.mark.parametrize("timeout", TIMEOUTS) @@ -366,7 +343,7 @@ def _assert_cancelled_cleanly( ids=["status", "get_log"], ) @patch("superset.db_engine_specs.impala.time.sleep") -def test_handle_cursor_cancels_on_a_clean_transport_after_interrupted_reply( +def test_handle_cursor_abandons_after_interrupted_reply( sleep_mock: Mock, interrupted: str, state: int, @@ -388,7 +365,7 @@ def test_handle_cursor_cancels_on_a_clean_transport_after_interrupted_reply( spec.handle_cursor(cursor, query) assert raised.value is timeout - _assert_cancelled_cleanly(mode, handler, interrupted, caplog) + _assert_abandoned(handler, interrupted) @pytest.mark.parametrize("timeout", TIMEOUTS) @@ -400,7 +377,7 @@ def test_handle_cursor_cancels_on_a_clean_transport_after_interrupted_reply( ], ids=["status", "rows"], ) -def test_fetch_data_cancels_on_a_clean_transport_after_interrupted_reply( +def test_fetch_data_abandons_after_interrupted_reply( interrupted: str, state: int, timeout: Exception, @@ -419,7 +396,7 @@ def test_fetch_data_cancels_on_a_clean_transport_after_interrupted_reply( spec.fetch_data(cursor) assert raised.value is timeout - _assert_cancelled_cleanly(mode, handler, interrupted, caplog) + _assert_abandoned(handler, interrupted) @pytest.mark.parametrize( @@ -503,8 +480,10 @@ def test_row_limited_fetch_sends_nothing_after_a_skipped_cancel( # The first batch fills the budget, so the probe fetches again. assert len(spec.fetch_data(limited)) == 500 _interrupt_after_header(cursor, "FetchResults", timeout) - # The probe reports a failed read as a possibly partial result. - assert limited.check_truncated(spec) is True + # A time limit must propagate, not become a partial result. + with pytest.raises(type(timeout)) as raised: + limited.check_truncated(spec) + assert raised.value is timeout cursor.__del__() assert cursor._closed is True @@ -544,7 +523,7 @@ def test_row_limited_fetch_sends_nothing_after_a_skipped_cancel( ], ) @patch("superset.db_engine_specs.impala.time.sleep") -def test_time_limit_during_cleanup_abandons_the_cursor( +def test_time_limit_prevents_cleanup_rpcs_and_second_reply_timeout( sleep_mock: Mock, entry: str, interrupted: str, @@ -587,7 +566,7 @@ def run() -> None: assert raised.value is timeout assert cursor._closed is True - assert handler.requests[-1] == cleanup_reply, handler.requests + assert handler.requests[-1] == interrupted, handler.requests @pytest.mark.parametrize( @@ -662,3 +641,173 @@ def test_reset_transport_refuses_an_unknown_transport() -> None: assert spec._reset_transport(cursor) is False transport.close.assert_not_called() transport.open.assert_not_called() + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize("interrupted", ["ExecuteStatement", "GetResultSetMetadata"]) +def test_execute_and_metadata_abandon_before_driver_cleanup( + interrupted: str, + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """Every RPC, including lazy metadata, disables all subsequent cleanup RPCs.""" + _, handler, conn = hs2_connection + cursor = conn.cursor() + if interrupted != "ExecuteStatement": + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + _interrupt_after_header(cursor, interrupted, timeout) + + def run() -> None: + """Issue the driver call that will be interrupted.""" + if interrupted == "ExecuteStatement": + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + else: + assert cursor.description + + with pytest.raises(type(timeout)) as raised, closing(cursor): + run() + cursor.__del__() + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == interrupted, handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "cleanup_reply", ["CancelOperation", "CloseOperation", "CloseSession"] +) +def test_cleanup_rpc_timeout_abandons_before_driver_handlers( + cleanup_reply: str, + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """A limit inside any cleanup reply must precede driver-internal cleanup.""" + _, handler, conn = hs2_connection + cursor = conn.cursor() + cursor.execute_async("SELECT 1") + _interrupt_after_header(cursor, cleanup_reply, timeout) + query = Mock(id=1, extra={QUERY_EARLY_CANCEL_KEY: True}, progress=0) + with patch("superset.db_engine_specs.impala.db"): + with pytest.raises(type(timeout)) as raised, closing(cursor): + spec.handle_cursor(cursor, query) + cursor.__del__() + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == cleanup_reply, handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@pytest.mark.parametrize( + "drop_failure", [OSError("drop failed"), SoftTimeLimitExceeded("drop")] +) +def test_second_limit_during_local_drop_preserves_execute_timeout( + drop_failure: Exception, + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """Even a second limit during local cleanup must not replace the first.""" + _, handler, conn = hs2_connection + cursor = conn.cursor() + _interrupt_after_header(cursor, "ExecuteStatement", timeout) + transport = cursor.session.service.client._iprot.trans + close = transport.close + attempts = 0 + + def interrupted_close() -> None: + nonlocal attempts + attempts += 1 + if attempts == 1: + raise drop_failure + close() + + with patch.object(transport, "close", side_effect=interrupted_close): + with pytest.raises(type(timeout)) as raised, closing(cursor): + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + cursor.__del__() + conn.close() + assert raised.value is timeout + assert attempts >= 2 + assert handler.requests[-1] == "ExecuteStatement", handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_connection_factory_guards_open_session_and_raw_cursor_calls( + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """The connection scope covers OpenSession before a cursor can be returned.""" + _, handler, conn = hs2_connection + spec.prepare_connection(conn) + _interrupt_replies(conn.service.client._iprot, [("OpenSession", timeout)]) + with pytest.raises(type(timeout)) as raised: + conn.cursor() + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == "OpenSession", handler.requests + + +def test_sigalrm_timeout_is_preserved_during_execute( + hs2_connection: tuple[str, FakeHiveServer2, Any], app_context: None +) -> None: + """Exercise utils.timeout's actual signal handler while a reply is unread.""" + import signal + + from superset.utils import core as utils + + _, handler, conn = hs2_connection + cursor = conn.cursor() + protocol = cursor.session.service.client._iprot + read = protocol.readMessageBegin + original: list[BaseException] = [] + + def interrupted_read() -> Any: + """Deliver SIGALRM after consuming the ExecuteStatement reply header.""" + header = read() + if header[0] == "ExecuteStatement": + try: + signal.raise_signal(signal.SIGALRM) + except SupersetTimeoutException as ex: + original.append(ex) + raise + return header + + protocol.readMessageBegin = interrupted_read + + def run() -> None: + """Close the real cursor after leaving the active timeout context.""" + with closing(cursor), utils.timeout(seconds=30): + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + + with pytest.raises(SupersetTimeoutException) as raised: + run() + conn.close() + cursor.__del__() + assert raised.value is original[0] + assert handler.requests[-1] == "ExecuteStatement", handler.requests + + +def test_timer_timeout_abandons_before_keyboard_interrupt_conversion( + hs2_connection: tuple[str, FakeHiveServer2, Any], app_context: None +) -> None: + """Windows converts KeyboardInterrupt only after inner cursor cleanup.""" + from superset.utils import core as utils + + _, handler, conn = hs2_connection + cursor = conn.cursor() + interrupt = KeyboardInterrupt() + _interrupt_after_header(cursor, "ExecuteStatement", interrupt) + + def run() -> None: + """Use the timer context without waiting for its background thread.""" + with utils.TimerTimeout(seconds=30), closing(cursor): + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + + with pytest.raises(SupersetTimeoutException, match="Timeout"): + run() + conn.close() + cursor.__del__() + assert handler.requests[-1] == "ExecuteStatement", handler.requests From 03bea0167cc3803399bbf18ed759cb84892ddf1d Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 17:45:02 +0000 Subject: [PATCH 12/14] fix(sqllab): guard SQLAlchemy HiveServer2 connection lifetimes --- superset/db_engine_specs/hive.py | 8 ++ superset/db_engine_specs/impala.py | 8 ++ superset/sql/execution/executor.py | 1 + superset/utils/hive.py | 15 ++- .../db_engine_specs/test_hive_transport.py | 103 ++++++++++++++++-- .../unit_tests/db_engine_specs/test_impala.py | 4 +- .../db_engine_specs/test_impala_pending.py | 38 ++++--- .../db_engine_specs/test_impala_transport.py | 50 ++++++++- 8 files changed, 192 insertions(+), 35 deletions(-) diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 261f6e844896..dc3a2b52dd0f 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -32,6 +32,7 @@ import pyarrow.parquet as pq from flask import current_app as app, g from sqlalchemy import Column, false, text, true, types +from sqlalchemy.engine import Engine from sqlalchemy.engine.interfaces import Dialect from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.url import URL @@ -54,6 +55,7 @@ from superset.utils.hive import ( guard_cursor, guard_cursor_factory, + guard_engine_connections, real_cursor, ThriftClientGuard, ) @@ -220,6 +222,12 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. except pyhive.exc.ProgrammingError: return [] + @classmethod + def register_engine_events(cls, engine: Engine) -> None: + """Guard raw connections used outside SQL Lab, including dialect pings.""" + super().register_engine_events(engine) + guard_engine_connections(engine, cls.prepare_connection) + @classmethod def prepare_connection(cls, connection: Any) -> None: """Guard every RPC on the PyHive connection's shared Thrift client.""" diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 98a43e31bb25..49f4a3cc0718 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -26,6 +26,7 @@ from flask import current_app as app from sqlalchemy import text, types +from sqlalchemy.engine import Engine from sqlalchemy.engine.reflection import Inspector from superset import db @@ -40,6 +41,7 @@ from superset.utils.hive import ( guard_cursor, guard_cursor_factory, + guard_engine_connections, real_cursor, ThriftClientGuard, ) @@ -168,6 +170,12 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. raise cls.get_dbapi_mapped_exception(ex) from ex return super().fetch_data(cursor, limit) + @classmethod + def register_engine_events(cls, engine: Engine) -> None: + """Guard raw connections used outside SQL Lab, including dialect pings.""" + super().register_engine_events(engine) + guard_engine_connections(engine, cls.prepare_connection) + @classmethod def prepare_connection(cls, connection: Any) -> None: """Guard OpenSession and all cursors sharing this connection's client.""" diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py index b6ace98a0ab0..8249fbc13c5a 100644 --- a/superset/sql/execution/executor.py +++ b/superset/sql/execution/executor.py @@ -235,6 +235,7 @@ def check_truncated(self, db_engine_spec: type[BaseEngineSpec]) -> bool: The rows within the budget are already materialized, so a failed probe reports the result as possibly partial instead of discarding them. + Time limits must propagate instead of becoming a partial result. """ if self._remaining != 0: return False diff --git a/superset/utils/hive.py b/superset/utils/hive.py index 5f7158bc7423..5d290b57c857 100644 --- a/superset/utils/hive.py +++ b/superset/utils/hive.py @@ -25,6 +25,9 @@ from functools import wraps from typing import Any, cast, TypeVar +from sqlalchemy import event +from sqlalchemy.engine import Engine + from superset.db_engine_specs.base import TIME_LIMIT_EXCEPTIONS logger = logging.getLogger(__name__) @@ -63,7 +66,7 @@ def discard(_ref: weakref.ReferenceType[Any]) -> None: def __getattr__(self, name: str) -> Any: """Forward attributes and wrap all callable driver-client entry points.""" value = getattr(self._client, name) - if not callable(value): + if name.startswith("_") or not callable(value): return value @wraps(value) @@ -123,3 +126,13 @@ def cursor(*args: Any, **kwargs: Any) -> Any: connection.cursor = cursor connection._superset_cursor_guard = True # pylint: disable=protected-access + + +def guard_engine_connections(engine: Engine, prepare: Callable[[Any], None]) -> None: + """Guard DB-API connections before dialect initialization and direct SQL use.""" + + def on_connect(connection: Any, _record: Any) -> None: + """Prepare a new connection before other SQLAlchemy connect listeners.""" + prepare(connection) + + event.listen(engine, "connect", on_connect, insert=True) diff --git a/tests/unit_tests/db_engine_specs/test_hive_transport.py b/tests/unit_tests/db_engine_specs/test_hive_transport.py index 902392769b40..ab1c2f821d2d 100644 --- a/tests/unit_tests/db_engine_specs/test_hive_transport.py +++ b/tests/unit_tests/db_engine_specs/test_hive_transport.py @@ -45,9 +45,10 @@ "tests.unit_tests.db_engine_specs.test_impala_transport" ) TTransportModule = pytest.importorskip("thrift.transport.TTransport") +reset_timeout_tracebacks = transport_tests.reset_timeout_tracebacks -@pytest.fixture(params=["NOSASL", "NONE"], ids=["binary", "sasl"]) +@pytest.fixture(params=["NOSASL", "NONE", "HTTP"], ids=["binary", "sasl", "http"]) def hive_connection( request: pytest.FixtureRequest, ) -> Iterator[tuple[Any, Any]]: @@ -55,6 +56,31 @@ def hive_connection( from pyhive import hive handler = transport_tests.FakeHiveServer2() + if request.param == "HTTP": + from thrift.transport.THttpClient import THttpClient + + server = transport_tests.THttpServer.THttpServer( + transport_tests.TCLIService.Processor(handler), + ("127.0.0.1", 0), + transport_tests._RecordingProtocolFactory(handler), + ) + thread = threading.Thread(target=server.serve, daemon=True) + thread.start() + transport = THttpClient( + f"http://127.0.0.1:{server.httpd.server_address[1]}/cliservice" + ) + transport.setTimeout(2000) + conn = hive.connect(thrift_transport=transport, username="user") + try: + yield handler, conn + finally: + if transport.isOpen(): + transport.close() + server.httpd.shutdown() + server.httpd.server_close() + thread.join(timeout=2) + return + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) listener.bind(("127.0.0.1", 0)) listener.listen() @@ -156,17 +182,29 @@ def test_failed_transport_drop_preserves_timeout_and_disables_session_close( handler, conn = hive_connection cursor = conn.cursor() cursor.execute("SELECT 1", async_=True) - transport_tests._interrupt_replies( - conn._client._iprot, [("GetOperationStatus", timeout)] - ) + protocol = conn._client._iprot + read = protocol.readMessageBegin + reply_interrupted = False + + def interrupted_read() -> Any: + """Arm the local-drop failure only after the reply has been interrupted.""" + nonlocal reply_interrupted + header = read() + if header[0] == "GetOperationStatus": + reply_interrupted = True + raise timeout + return header + + protocol.readMessageBegin = interrupted_read close = conn._transport.close attempts = 0 def interrupted_close() -> None: nonlocal attempts - attempts += 1 - if attempts == 1: - raise drop_failure + if reply_interrupted: + attempts += 1 + if attempts == 1: + raise drop_failure close() with patch.object(conn._transport, "close", side_effect=interrupted_close): @@ -204,19 +242,60 @@ def test_connection_close_sends_no_rpc_after_an_interrupted_reply( @pytest.mark.parametrize("timeout", TIMEOUTS) -def test_execute_abandons_before_cursor_and_connection_cleanup( +@pytest.mark.parametrize("interrupted", ["ExecuteStatement", "GetResultSetMetadata"]) +def test_execute_and_metadata_abandon_before_cursor_and_connection_cleanup( + interrupted: str, timeout: Exception, hive_connection: tuple[Any, Any], app_context: None, ) -> None: - """An interrupted ExecuteStatement must not be followed by CloseSession.""" + """Execution and lazy metadata must be abandoned before any cleanup RPC.""" handler, conn = hive_connection cursor = conn.cursor() + if interrupted != "ExecuteStatement": + HiveEngineSpec.execute(cursor, "SELECT 1", None, async_=True) # type: ignore[arg-type] + transport_tests._interrupt_replies(conn._client._iprot, [(interrupted, timeout)]) + + def run() -> None: + """Read a reply, leaving its body unread when the limit fires.""" + if interrupted == "ExecuteStatement": + HiveEngineSpec.execute(cursor, "SELECT 1", None, async_=True) # type: ignore[arg-type] + else: + assert cursor.description + + with pytest.raises(type(timeout)) as raised, closing(cursor): + run() + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == interrupted, handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_sqlalchemy_ping_abandons_before_dialect_cursor_cleanup( + timeout: Exception, + hive_connection: tuple[Any, Any], + app_context: None, +) -> None: + """Connection validation must not reuse an interrupted transport either.""" + from sqlalchemy import create_engine + from sqlalchemy.pool import NullPool + + handler, conn = hive_connection + engine = create_engine("hive://", creator=lambda: conn, poolclass=NullPool) + HiveEngineSpec.register_engine_events(engine) transport_tests._interrupt_replies( conn._client._iprot, [("ExecuteStatement", timeout)] ) - with pytest.raises(type(timeout)) as raised, closing(cursor): - HiveEngineSpec.execute(cursor, "SELECT 1", None, async_=True) # type: ignore[arg-type] - conn.close() + + def run() -> None: + """Exercise the raw-connection ping, without calling spec.execute().""" + with closing(engine.raw_connection()) as raw: + engine.dialect.do_ping(raw) + + try: + with pytest.raises(type(timeout)) as raised: + run() + finally: + engine.dispose() assert raised.value is timeout assert handler.requests[-1] == "ExecuteStatement", handler.requests diff --git a/tests/unit_tests/db_engine_specs/test_impala.py b/tests/unit_tests/db_engine_specs/test_impala.py index 296431eb0a23..6547cb667c63 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -411,7 +411,7 @@ def test_fetch_data_abandons_the_operation_on_time_limit( "cancel_timeout_type", [SoftTimeLimitExceeded, SupersetTimeoutException] ) @patch("superset.db_engine_specs.impala.time.sleep") -def test_fetch_data_propagates_time_limit_during_cancel( +def test_fetch_data_never_enters_cancellation_after_time_limit( sleep_mock: Mock, cancel_timeout_type: type[Exception], app_context: None, @@ -663,7 +663,7 @@ def test_time_limit_during_transport_reset_propagates( ], ids=["async", "sync"], ) -def test_time_limit_during_cleanup_abandons_the_cursor( +def test_fetch_time_limit_never_enters_rpc_cleanup( cleanup_rpc: str, second_timeout: Exception, clean_transport: Mock, diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 6af03012eb0e..eedb68b2e313 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_pending.py +++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py @@ -58,7 +58,7 @@ def rejected_reply(mocker: MockerFixture) -> type[Exception]: @pytest.fixture(autouse=True) def clean_transport() -> Iterator[Mock]: """ - Report the transport reset as successful for mock cursors. + Stub local transport drops for mock cursors. Mock cursors have no Thrift transport; resets over real impyla transports are covered in ``test_impala_transport.py``. @@ -348,7 +348,7 @@ def test_pending_operation_polls_with_backoff() -> None: @pytest.mark.parametrize("state", ["PENDING_STATE", "RUNNING_STATE"]) -def test_soft_time_limit_cancels_the_operation(state: str) -> None: +def test_soft_time_limit_abandons_the_operation(state: str) -> None: """ The soft time limit firing while the loop polls cancels the operation and propagates, so SQL Lab marks the query timed out instead of waiting on it. @@ -369,9 +369,11 @@ def test_soft_time_limit_cancels_the_operation(state: str) -> None: ): with pytest.raises(SoftTimeLimitExceeded): ImpalaEngineSpec.handle_cursor(cursor, query) - cursor.cancel_operation.assert_called_once_with() - cursor.close_operation.assert_called_once_with() - cursor.close.assert_called_once_with() + cursor.cancel_operation.assert_not_called() + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() + + assert cursor._closed is True @pytest.mark.parametrize("method", ["cancel_operation", "close_operation", "close"]) @@ -470,7 +472,7 @@ def _sync_timeout() -> SupersetTimeoutException: @pytest.mark.parametrize("state", ["PENDING_STATE", "RUNNING_STATE"]) -def test_sync_timeout_cancels_the_operation(state: str) -> None: +def test_sync_timeout_abandons_the_operation(state: str) -> None: """ The synchronous ``SQLLAB_TIMEOUT`` firing while the loop polls must cancel the operation and propagate, rather than being swallowed so the Run keeps waiting @@ -490,15 +492,17 @@ def test_sync_timeout_cancels_the_operation(state: str) -> None: ): with pytest.raises(SupersetTimeoutException): ImpalaEngineSpec.handle_cursor(cursor, query) - cursor.cancel_operation.assert_called_once_with() - cursor.close_operation.assert_called_once_with() - cursor.close.assert_called_once_with() + cursor.cancel_operation.assert_not_called() + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() + + assert cursor._closed is True @pytest.mark.parametrize( "timeout", [SoftTimeLimitExceeded(), _sync_timeout()], ids=["async", "sync"] ) -def test_timeout_during_progress_log_cancels_the_operation( +def test_timeout_during_progress_log_abandons_the_operation( timeout: Exception, ) -> None: """ @@ -522,9 +526,11 @@ def test_timeout_during_progress_log_cancels_the_operation( with pytest.raises(type(timeout)): ImpalaEngineSpec.handle_cursor(cursor, query) sleep.assert_not_called() - cursor.cancel_operation.assert_called_once_with() - cursor.close_operation.assert_called_once_with() - cursor.close.assert_called_once_with() + cursor.cancel_operation.assert_not_called() + cursor.close_operation.assert_not_called() + cursor.close.assert_not_called() + + assert cursor._closed is True def test_status_failure_still_ends_polling_without_raising() -> None: @@ -538,7 +544,7 @@ def test_status_failure_still_ends_polling_without_raising() -> None: cursor.cancel_operation.assert_not_called() -def test_sync_timeout_fails_the_run_and_cancels_the_operation( +def test_sync_timeout_fails_the_run_and_abandons_the_operation( mocker: MockerFixture, app_context: None ) -> None: """A synchronous timeout during polling ends the Run with the timeout message.""" @@ -567,5 +573,7 @@ def test_sync_timeout_fails_the_run_and_cancels_the_operation( with pytest.raises(SqlLabException, match="exceeded the 30 seconds timeout"): execute_query(query, cursor) - cursor.cancel_operation.assert_called_once_with() + cursor.cancel_operation.assert_not_called() cursor.fetchall.assert_not_called() + + assert cursor._closed is True diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py index a597c24e1613..9e728defed5c 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_transport.py +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -25,6 +25,7 @@ from __future__ import annotations +import signal import socket import struct import threading @@ -325,6 +326,17 @@ def _interrupt_after_header(cursor: Any, reply: str, timeout: BaseException) -> ] +@pytest.fixture(autouse=True) +def reset_timeout_tracebacks(request: pytest.FixtureRequest) -> None: + """Keep parametrized exception instances from retaining previous cases.""" + if (callspec := getattr(request.node, "callspec", None)) is not None: + for value in callspec.params.values(): + if isinstance(value, BaseException): + value.__traceback__ = None + value.__cause__ = None + value.__context__ = None + + def _assert_abandoned( handler: FakeHiveServer2, interrupted: str, @@ -534,10 +546,10 @@ def test_time_limit_prevents_cleanup_rpcs_and_second_reply_timeout( app_context: None, ) -> None: """ - A second limit inside a cleanup reply leaves that reply unread too. + A first limit prevents entering cleanup, including an armed second limit. - Nothing may be sent after it: not by the cancellation helper, the - caller's ``closing()``, nor the cursor's finalizer. + Nothing may be sent after the first interrupted reply: not by cancellation, + the caller's ``closing()``, nor the cursor's finalizer. """ _, handler, conn = hs2_connection if entry == "fetch_data": @@ -750,12 +762,11 @@ def test_connection_factory_guards_open_session_and_raw_cursor_calls( assert handler.requests[-1] == "OpenSession", handler.requests +@pytest.mark.skipif(not hasattr(signal, "SIGALRM"), reason="SIGALRM is unavailable") def test_sigalrm_timeout_is_preserved_during_execute( hs2_connection: tuple[str, FakeHiveServer2, Any], app_context: None ) -> None: """Exercise utils.timeout's actual signal handler while a reply is unread.""" - import signal - from superset.utils import core as utils _, handler, conn = hs2_connection @@ -811,3 +822,32 @@ def run() -> None: conn.close() cursor.__del__() assert handler.requests[-1] == "ExecuteStatement", handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_sqlalchemy_ping_abandons_before_dialect_cursor_cleanup( + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """Direct engine users get the guard before the dialect creates a cursor.""" + from sqlalchemy import create_engine + from sqlalchemy.pool import NullPool + + _, handler, conn = hs2_connection + engine = create_engine("impala://", creator=lambda: conn, poolclass=NullPool) + spec.register_engine_events(engine) + _interrupt_replies(conn.service.client._iprot, [("ExecuteStatement", timeout)]) + + def run() -> None: + """Exercise the same raw-connection ping used by connection validation.""" + with closing(engine.raw_connection()) as raw: + engine.dialect.do_ping(raw) + + try: + with pytest.raises(type(timeout)) as raised: + run() + finally: + engine.dispose() + assert raised.value is timeout + assert handler.requests[-1] == "ExecuteStatement", handler.requests From bb609de6337fcdbff1bd5aa74301782a61b06a19 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 17:50:14 +0000 Subject: [PATCH 13/14] fix(impala): prevent driver cleanup from reopening abandoned transports --- superset/db_engine_specs/impala.py | 19 +++++++++------ superset/utils/hive.py | 8 +++++-- .../db_engine_specs/test_impala_transport.py | 24 +++++++++++++++++++ 3 files changed, 42 insertions(+), 9 deletions(-) diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 49f4a3cc0718..a40e56261331 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -189,6 +189,9 @@ def prepare_connection(cls, connection: Any) -> None: def _abandon_service(service: Any) -> None: """Make connection close local-only and non-fatal after interruption.""" transport = service.client._iprot.trans + # impyla opens a closed transport before looking up the RPC method. + # Reject that open too, including inside the driver's own close(). + transport.open = service.client.check_active def close() -> None: try: @@ -253,10 +256,9 @@ def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: connection, while the buffering layers above it keep their read and write buffers. Returns whether every layer was recognised and reset. - With ``reconnect`` the transport is opened again. Without it a socket - is left closed until impyla's next RPC opens it; an HTTP client is - still reopened, which creates a connection object without network I/O, - because closing it a second time raises. + With ``reconnect`` the transport is opened again. Without it every + transport remains closed; an abandoned client's open guard also stops + driver-internal cleanup from reconnecting before attempting an RPC. """ # pylint: disable=import-outside-toplevel from io import BytesIO @@ -268,7 +270,7 @@ def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: if not reconnect: transport.close() return False - buffers, over_socket = stack + buffers, _ = stack # Closing drops the socket together with any reply bytes still in it. try: @@ -281,7 +283,7 @@ def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: raise for buffered_layer, attribute in buffers: setattr(buffered_layer, attribute, BytesIO()) - if reconnect or not over_socket: + if reconnect: transport.open() return True @@ -388,7 +390,10 @@ def _cancel_operation(cls, cursor: Any, query_id: int | None = None) -> None: except Exception as ex: # pylint: disable=broad-except logger.warning("Query %s: %s() failed", query_id, name) if not cls._is_error_reply(ex): - cls._abandon_cursor(cursor, query_id) + if on_error := cls.prepare_cursor(cursor): + on_error(ex) + else: + cls._abandon_cursor(cursor, query_id) return @classmethod diff --git a/superset/utils/hive.py b/superset/utils/hive.py index 5d290b57c857..8d6280209a08 100644 --- a/superset/utils/hive.py +++ b/superset/utils/hive.py @@ -63,6 +63,11 @@ def discard(_ref: weakref.ReferenceType[Any]) -> None: self._targets[key] = (weakref.ref(target, discard), abandon) + def check_active(self) -> None: + """Reject transport reopen and RPC attempts after abandonment.""" + if self._exception is not None: + raise self._exception + def __getattr__(self, name: str) -> Any: """Forward attributes and wrap all callable driver-client entry points.""" value = getattr(self._client, name) @@ -71,8 +76,7 @@ def __getattr__(self, name: str) -> Any: @wraps(value) def call(*args: Any, **kwargs: Any) -> Any: - if self._exception is not None: - raise self._exception + self.check_active() try: return value(*args, **kwargs) except TIME_LIMIT_EXCEPTIONS as ex: diff --git a/tests/unit_tests/db_engine_specs/test_impala_transport.py b/tests/unit_tests/db_engine_specs/test_impala_transport.py index 9e728defed5c..6d4582272554 100644 --- a/tests/unit_tests/db_engine_specs/test_impala_transport.py +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -851,3 +851,27 @@ def run() -> None: engine.dispose() assert raised.value is timeout assert handler.requests[-1] == "ExecuteStatement", handler.requests + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +def test_driver_close_timeout_prevents_transport_reopen( + timeout: Exception, + hs2_connection: tuple[str, FakeHiveServer2, Any], + app_context: None, +) -> None: + """Driver-internal CloseSession must not even reopen the abandoned socket.""" + _, handler, conn = hs2_connection + cursor = conn.cursor() + spec.execute(cursor, "SELECT 1", None) # type: ignore[arg-type] + _interrupt_after_header(cursor, "CloseOperation", timeout) + transport = cursor.session.service.client._iprot.trans + with patch.object( + transport, "open", side_effect=AssertionError("transport reopened") + ) as reopen: + with pytest.raises(type(timeout)) as raised: + cursor.close() + reopen.assert_not_called() + cursor.__del__() + conn.close() + assert raised.value is timeout + assert handler.requests[-1] == "CloseOperation", handler.requests From 8c8604ec5640098c383131e810d4caf7bc5a96c4 Mon Sep 17 00:00:00 2001 From: Amin Ghadersohi Date: Sat, 31 Oct 2026 20:21:16 +0000 Subject: [PATCH 14/14] test(sql-lab): cover time limits during truncation probes --- .../unit_tests/sql/execution/test_executor.py | 57 +++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py index df3f9062f0d5..5fafce0fd4ff 100644 --- a/tests/unit_tests/sql/execution/test_executor.py +++ b/tests/unit_tests/sql/execution/test_executor.py @@ -46,6 +46,7 @@ QueryStatus, ) +from superset.db_engine_specs.base import TIME_LIMIT_EXCEPTIONS from superset.models.core import Database from superset.sql.parse import LimitMethod, SQLScript @@ -3486,6 +3487,62 @@ def test_limited_cursor_truncation_probe(row_count: int, limit: int) -> None: assert limited.check_truncated(BaseEngineSpec) is (row_count > limit) +@pytest.mark.parametrize("timeout_type", TIME_LIMIT_EXCEPTIONS) +def test_execute_truncation_probe_propagates_time_limits( + mocker: MockerFixture, + mock_database: MagicMock, + mock_query: MagicMock, + app_context: None, + caplog: pytest.LogCaptureFixture, + timeout_type: type[BaseException], +) -> None: + """A probe timeout propagates unchanged instead of returning partial rows.""" + from superset.db_engine_specs.base import BaseEngineSpec + from superset.errors import ErrorLevel, SupersetErrorType + from superset.exceptions import SupersetTimeoutException + from superset.sql.execution.executor import execute_sql_with_cursor + + timeout = ( + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="Truncation probe timed out", + level=ErrorLevel.ERROR, + ) + if timeout_type is SupersetTimeoutException + else timeout_type() + ) + mock_database.db_engine_spec = BaseEngineSpec + mock_query.limit = 2 + cursor = create_mock_cursor(["n"]) + cursor.fetchmany.return_value = [(1,), (2,)] + result_set = mocker.patch("superset.result_set.SupersetResultSet") + + def fetch_rows(fetch_cursor: Any) -> list[tuple[int, ...]]: + """Consume the row budget, then interrupt the next engine fetch.""" + rows = fetch_cursor.fetchall() + fetch_data.side_effect = timeout + return rows + + fetch_data = mocker.patch.object( + BaseEngineSpec, "fetch_data", side_effect=fetch_rows + ) + + with pytest.raises(timeout_type) as raised: + execute_sql_with_cursor( + mock_database, + cursor, + ["SELECT n FROM t"], + mock_query, + execute_fn=MagicMock(), + ) + + assert raised.value is timeout + assert fetch_data.call_count == 2 + assert fetch_data.call_args.args[0]._remaining == 1 + assert "Truncation probe failed" not in caplog.text + result_set.assert_not_called() + + @pytest.mark.parametrize("eof", [True, False]) def test_execute_truncation_probe_engine_errors( mocker: MockerFixture,