diff --git a/docs/admin_docs/configuration/configuring-superset.mdx b/docs/admin_docs/configuration/configuring-superset.mdx index 97cfb3f343e7..c15be805d124 100644 --- a/docs/admin_docs/configuration/configuring-superset.mdx +++ b/docs/admin_docs/configuration/configuring-superset.mdx @@ -833,7 +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. +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 7eadb4a94d56..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 @@ -41,6 +43,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 +80,7 @@ OAuth2TokenRefreshError, SupersetGenericDBErrorException, SupersetParseError, + SupersetTimeoutException, ) from superset.key_value.types import JsonKeyValueCodec, KeyValueResource from superset.sql.parse import ( @@ -134,6 +138,19 @@ # that we know this is a necessary evil we create an alias, and catch it instead. GenericDBException = Exception +# SQL Lab's Celery soft limit for asynchronous runs and ``SQLLAB_TIMEOUT`` +# SIGALRM for main-thread synchronous runs reach the cancellation handlers. +# The alarm is main-thread only: threaded synchronous requests are not covered +# by this change. Each limit fires once, so swallowing it leaves the query running +# with nothing left to stop it. +# 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]: result_set_columns: list[ResultSetColumnType] = [] @@ -2062,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 b85868ea7735..dc3a2b52dd0f 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 @@ -31,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 @@ -39,13 +41,24 @@ 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 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, + guard_engine_connections, + real_cursor, + ThriftClientGuard, +) if TYPE_CHECKING: from superset.models.core import Database @@ -193,6 +206,7 @@ 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 @@ -208,6 +222,52 @@ 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.""" + 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_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 + cls._abandon_connection(cursor._connection) # pylint: disable=protected-access + @classmethod def df_to_sql( cls, @@ -429,10 +489,16 @@ 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 + @guard_cursor + def handle_cursor(cls, cursor: Any, query: Query) -> None: + """Updates progress information""" + cls._poll_cursor(cursor, query) + + @classmethod + def _poll_cursor( # pylint: disable=too-many-locals # noqa: C901 cls, cursor: Any, query: Query ) -> None: - """Updates progress information""" + """Poll the operation until it finishes, recording its progress.""" # pylint: disable=import-outside-toplevel from pyhive import hive @@ -458,7 +524,9 @@ def handle_cursor( # pylint: disable=too-many-locals # noqa: C901 try: logs = cursor.fetch_logs() log = "\n".join(logs) if logs else "" - except Exception: # pylint: disable=broad-except + except TIME_LIMIT_EXCEPTIONS: + raise + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.warning("Call to GetLog() failed") log = "" @@ -528,7 +596,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: @@ -629,14 +697,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 e5134e21b947..a40e56261331 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -20,19 +20,31 @@ import logging import re import time +from collections.abc import Callable 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 import Engine 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.hive import ( + guard_cursor, + guard_cursor_factory, + guard_engine_connections, + real_cursor, + ThriftClientGuard, +) from superset.utils.network import get_ssrf_safe_requester, is_safe_host if TYPE_CHECKING: @@ -121,6 +133,7 @@ def has_implicit_cancel(cls) -> bool: return False @classmethod + @guard_cursor def execute( cls, cursor: Any, @@ -130,10 +143,13 @@ 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.""" if callable(getattr(cursor, "execute_async", None)): @@ -143,7 +159,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") time.sleep(0.1) if cursor.execution_failed(): @@ -155,31 +171,233 @@ def fetch_data(cls, cursor: Any, limit: int | None = None) -> list[tuple[Any, .. return super().fetch_data(cursor, limit) @classmethod - def _cancel_operation(cls, cursor: Any, query_id: int) -> None: - """Cancel the live operation and release its handles.""" + 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.""" + 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 + # 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: + 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 + # 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: - cursor.cancel_operation() - except SoftTimeLimitExceeded: - 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. + cls._reset_transport(cursor, reconnect=False) + 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) + + @classmethod + def _reset_transport(cls, cursor: Any, reconnect: bool = True) -> bool: + """ + 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 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 + + # 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, _ = stack + + # Closing drops the socket together with any reply bytes still in it. try: - cursor.close_operation() - except SoftTimeLimitExceeded: + transport.close() + except TIME_LIMIT_EXCEPTIONS: raise - except Exception: # pylint: disable=broad-except - logger.warning("Query %s: close_operation() failed", query_id) + except (Exception, KeyboardInterrupt): # 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: + 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 + + # (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: - cursor.close() - except SoftTimeLimitExceeded: - raise - except Exception: # pylint: disable=broad-except - logger.warning("Query %s: close() failed", query_id) + from thrift_sasl import TSaslClientTransport + + layouts.append( + ( + TSaslClientTransport, + ("_TSaslClientTransport__rbuf", "_TSaslClientTransport__wbuf"), + "_trans", + ) + ) + except ImportError: + pass + + 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 None + _, attributes, inner = layout + if not all(hasattr(layer, attribute) for attribute in attributes): + return None + buffers.extend((layer, attribute) for attribute in attributes) + if inner is None: + return buffers, False + layer = getattr(layer, inner, None) + return buffers, True + + @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: + 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.""" + # 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): + if on_error := cls.prepare_cursor(cursor): + on_error(ex) + else: + cls._abandon_cursor(cursor, query_id) + return @classmethod + @guard_cursor def handle_cursor(cls, cursor: Any, query: Query) -> None: """Stop query and updates progress information""" @@ -219,7 +437,9 @@ def handle_cursor(cls, cursor: Any, query: Query) -> None: if status == "RUNNING_STATE": try: log = cursor.get_log() or "" - except Exception: # pylint: disable=broad-except + except TIME_LIMIT_EXCEPTIONS: + raise + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.warning("Call to GetLog() failed") log = "" @@ -234,13 +454,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 - # coordinator while fetch_data() waits on it. - cls._cancel_operation(cursor, query_id) + except TIME_LIMIT_EXCEPTIONS: raise - except Exception: # pylint: disable=broad-except + except (Exception, KeyboardInterrupt): # pylint: disable=broad-except logger.debug("Call to status() failed ") return @@ -338,7 +554,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 1aee2fd441f7..8249fbc13c5a 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, @@ -205,15 +206,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 + def __setattr__(self, name: str, value: Any) -> None: + """ + Write driver state, such as a closed flag, to the real cursor. - @arraysize.setter - def arraysize(self, value: int) -> None: - """Preserve engine-specific cursor batch-size configuration.""" - self._cursor.arraysize = value + 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.""" @@ -232,11 +235,14 @@ 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 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", @@ -918,6 +924,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( @@ -1644,6 +1651,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..8d6280209a08 --- /dev/null +++ b/superset/utils/hive.py @@ -0,0 +1,142 @@ +# 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 sqlalchemy import event +from sqlalchemy.engine import Engine + +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 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) + if name.startswith("_") or not callable(value): + return value + + @wraps(value) + def call(*args: Any, **kwargs: Any) -> Any: + self.check_active() + 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 + + +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.py b/tests/unit_tests/db_engine_specs/test_hive.py index dc6f9b071e1b..36ea92e3236c 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,105 @@ 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 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) + 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=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" + 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..ab1c2f821d2d --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_hive_transport.py @@ -0,0 +1,301 @@ +# 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") +reset_timeout_tracebacks = transport_tests.reset_timeout_tracebacks + + +@pytest.fixture(params=["NOSASL", "NONE", "HTTP"], ids=["binary", "sasl", "http"]) +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() + 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() + 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) + 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 + if reply_interrupted: + 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 + + +@pytest.mark.parametrize("timeout", TIMEOUTS) +@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: + """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)] + ) + + 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 f8603e2f5410..6547cb667c63 100644 --- a/tests/unit_tests/db_engine_specs/test_impala.py +++ b/tests/unit_tests/db_engine_specs/test_impala.py @@ -15,16 +15,21 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Iterator from datetime import datetime -from typing import Optional -from unittest.mock import Mock, patch +from typing import Any, Optional +from unittest.mock import call, MagicMock, 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 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 @@ -182,10 +187,33 @@ 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]: + """ + 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``. + """ + 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.""" - mocker.patch.dict("sys.modules", {"impala.error": Mock(Error=_DriverError)}) + mocker.patch.dict( + "sys.modules", + { + "impala.error": Mock( + Error=_DriverError, HiveServer2Error=_RejectedReplyError + ) + }, + ) class _AsyncCursor: @@ -254,6 +282,53 @@ 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 rejected cancel must not mask the deadline or retain operation handles.""" + from flask import current_app + + cursor = Mock() + cursor.is_executing.return_value = True + 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"): + 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 + + +@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( @@ -304,3 +379,352 @@ 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_abandons_the_operation_on_time_limit( + sleep_mock: Mock, timeout: Exception, app_context: None, impala_errors: None +) -> None: + """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_not_called() + 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_never_enters_cancellation_after_time_limit( + sleep_mock: Mock, + cancel_timeout_type: type[Exception], + app_context: None, + impala_errors: None, +) -> None: + """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: + 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(SoftTimeLimitExceeded) as raised: + spec.fetch_data(cursor) + + assert raised.value is original_timeout + cursor.cancel_operation.assert_not_called() + 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_abandons_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_not_called() + + +@pytest.mark.parametrize("cancel_fails", [False, True]) +def test_fetch_data_drops_transport_without_cancelling_an_interrupted_fetch( + cancel_fails: bool, + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """Abandon an interrupted response without issuing any cleanup RPC.""" + 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(clean_transport, "reset_transport") + calls.attach_mock(cursor.cancel_operation, "cancel_operation") + if cancel_fails: + cursor.cancel_operation.side_effect = _RejectedReplyError("cancel rejected") + 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 == [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"]) +@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_drops_transport_without_cancelling_an_interrupted_poll( + interrupted_rpc: str, + cancel_fails: bool, + 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 = _RejectedReplyError("cancel rejected") + query = Mock(id=1, extra={}, status="running", progress=0) + calls = MagicMock() + 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") + + 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.reset_transport(cursor, reconnect=False), + ] + + +@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 + # 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() + # 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 limit while resetting must not hide the original time limit.""" + 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 + + 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 + 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_fetch_time_limit_never_enters_rpc_cleanup( + cleanup_rpc: str, + second_timeout: Exception, + clean_transport: Mock, + app_context: None, + impala_errors: None, +) -> None: + """An expired fetch never enters RPC cleanup, even if it would time out.""" + 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, 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._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_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py index 1e7d69b76e78..eedb68b2e313 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 @@ -29,11 +30,45 @@ 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 +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]: + """ + 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``. + """ + 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 @@ -222,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() @@ -311,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. @@ -332,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"]) @@ -362,15 +401,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) @@ -379,6 +421,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__) @@ -401,3 +460,120 @@ 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_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 + 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_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_abandons_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_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: + """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_abandons_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_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 new file mode 100644 index 000000000000..6d4582272554 --- /dev/null +++ b/tests/unit_tests/db_engine_specs/test_impala_transport.py @@ -0,0 +1,877 @@ +# 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 signal +import socket +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 + +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") +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") +TProtocol = pytest.importorskip("thrift.protocol.TProtocol") +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 + 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) + 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=self.has_more_rows, + 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=self.cancel_status) + + 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 _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: + 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 = _recording_protocol(transport, handler) + 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), + _RecordingProtocolFactory(handler), + ) + 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: + conn.close() + listener.shutdown(socket.SHUT_RDWR) + listener.close() + thread.join(timeout=2) + + +def _interrupt_replies( + protocol: Any, interrupts: list[tuple[str, BaseException]] +) -> 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 + pending = list(interrupts) + + def interrupted_read_message_begin() -> Any: + header = read_message_begin() + 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: BaseException) -> 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( + SupersetTimeoutException( + error_type=SupersetErrorType.BACKEND_TIMEOUT_ERROR, + message="The query exceeded the 30 seconds timeout.", + level=ErrorLevel.ERROR, + ), + id="sync", + ), +] + + +@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, +) -> None: + """The interrupted reply must be the last RPC, without reconnection.""" + assert handler.requests[-1] == interrupted, handler.requests + + +@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_abandons_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_abandoned(handler, interrupted) + + +@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_abandons_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_abandoned(handler, interrupted) + + +@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 :] == [] + + +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) + # 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 + 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_prevents_cleanup_rpcs_and_second_reply_timeout( + 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 first limit prevents entering cleanup, including an armed second limit. + + 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": + 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] == interrupted, 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() + transport = cursor.session.service.client._iprot.trans + + 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 + + +@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.""" + 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 + + +@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 + + +@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 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."""