aminghadersohi commented on code in PR #45845:
URL: https://github.com/apache/superset/pull/45845#discussion_r4450484381
##########
superset/db_engine_specs/impala.py:
##########
@@ -130,56 +143,311 @@ 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)):
- from impala.error import Error
+ """Wait for asynchronous operations using the public cursor API.
+
+ On a time limit, only cancel before abandoning locally: alerts/reports
+ have a short soft-to-hard limit gap. CancelOperation itself is not
+ bounded here; configure a driver network timeout or increase
+ ALERT_REPORTS_WORKING_TIME_OUT_LAG for Impala alerts/reports.
+ """
+ try:
+ if callable(getattr(cursor, "execute_async", None)):
+ from impala.error import Error
+
+ deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"]
+ try:
+ while cursor.is_executing():
+ if time.monotonic() >= deadline:
+ cls._cancel_operation(cursor)
+ raise TimeoutError(
+ "Timed out waiting for the Impala operation"
+ )
+ time.sleep(0.1)
+ if cursor.execution_failed():
+ # Fetching surfaces the driver's detailed asynchronous
error,
+ # including errors on statements that produce no
result set.
+ cursor.fetchall()
+ except Error as ex:
+ raise cls.get_dbapi_mapped_exception(ex) from ex
+ return super().fetch_data(cursor, limit)
+ except TIME_LIMIT_EXCEPTIONS as ex:
+ cls._cancel_operation_on_time_limit(cursor, fetch_time_limit=ex)
+ raise
+
+ @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)
- deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"]
+ @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:
- while cursor.is_executing():
- if time.monotonic() >= deadline:
- cursor.cancel_operation()
- raise TimeoutError("Timed out waiting for the Impala
operation")
- time.sleep(0.1)
- if cursor.execution_failed():
- # Fetching surfaces the driver's detailed asynchronous
error,
- # including errors on statements that produce no result
set.
- cursor.fetchall()
- except Error as ex:
- raise cls.get_dbapi_mapped_exception(ex) from ex
- return super().fetch_data(cursor, limit)
+ transport.close()
+ except (Exception, KeyboardInterrupt): # pylint:
disable=broad-except
+ logger.warning("Dropping the interrupted Impala transport
failed")
+
+ service.close = close
+ close()
@classmethod
- def _cancel_operation(cls, cursor: Any, query_id: int) -> None:
- """Cancel the live operation and release its handles."""
+ 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.on_time_limit
+
+ @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 cls.prepare_cursor(cursor) is not None:
+ guard = real_cursor(cursor).session.service.client
+ guard.abandon(ex)
+ else:
+ cls._abandon_cursor(cursor, query_id)
+ return
+
+ @classmethod
+ def _cancel_operation_on_time_limit(
+ cls,
+ cursor: Any,
+ query_id: int | None = None,
+ *,
+ fetch_time_limit: BaseException | None = None,
+ ) -> None:
+ """Cancel after a limit between RPCs, never on an interrupted
transport."""
+ cursor = real_cursor(cursor)
+ if (
+ getattr(cursor, "_superset_abandoned", False) is True
+ or getattr(cursor, "_closed", False) is True
+ ):
+ return
+ session = getattr(cursor, "session", None)
+ guard = session.service.client if session is not None else None
+ if isinstance(guard, ThriftClientGuard) and guard.interrupted:
+ return
+ if fetch_time_limit is not None:
+ try:
+ # The default resets state via CloseOperation, which can block.
+ cursor.cancel_operation(reset_state=False)
+ except (Exception, KeyboardInterrupt): # pylint:
disable=broad-except
+ logger.warning("Query %s: cancel_operation() failed", query_id)
+ finally:
+ if isinstance(guard, ThriftClientGuard):
+ guard.abandon(fetch_time_limit)
Review Comment:
Fixed in 7ea9cdbe1bddccd833982c5df5a84fef931cb4dd.
I reproduced this before changing production code. A size-1 SQLAlchemy
QueuePool with the real impyla driver reuses the abandoned DB-API connection;
the second checkout reaches `check_active` and raises the first query's
timeout. That initial reproduction failed in all 6 soft/synchronous
binary/SASL/HTTP cases.
The shared engine hook now invalidates abandoned pool records on checkin. It
checks both the latched guard and local-abandonment markers, so
cursor/connection drops that do not latch an exception are discarded too.
Healthy connections remain pooled. This operates on individual records; it does
not dispose the cached engine or reset the old guard. The configuration
documentation describes the pool-return behavior.
Added regression tests:
- `test_pooled_connection_is_replaced_after_abandonment`: 42 cases covering
Impala fetch-time, poll-time, cancel-failure, execute/metadata RPC
interruption, direct guard abandonment and local cursor abandonment.
- `test_hive_pool_replaces_abandoned_connections`: 42 cases covering Hive
poll/fetch/execute/cleanup RPC interruption, direct guard abandonment, local
cursor abandonment and local connection abandonment.
- `test_pool_reuses_healthy_impala_connections` and
`test_pool_reuses_healthy_hive_connections`: 6 healthy-reuse controls.
Before: retaining the final tests while restoring the three production
helper/spec files from 1cb704263a5b816912fa7b562f84b857ad140a38 reported `84
failed, 6 passed`. After: all 90 cases pass. Each abandonment regression
verifies a different DB-API connection is checked out, the subsequent query
sends ExecuteStatement successfully, and no further requests reach the
abandoned connection.
Pool-regression command:
```bash
python -m pytest -n auto --dist loadfile -q
tests/unit_tests/db_engine_specs/test_impala_transport.py
tests/unit_tests/db_engine_specs/test_hive_transport.py -k
'test_pooled_connection_is_replaced_after_abandonment or
test_hive_pool_replaces_abandoned_connections or test_pool_reuses_healthy_' -m
'not manual and not llm_compliance'
```
The scoped CI-flags pytest run, including Impala/Hive/Base, both transport
suites, Presto/Trino, executor and workflow tests, reported `1318 passed, 3
xfailed`, exit 0. Per-file transport counts are 310 Impala and 111 Hive passes,
zero skips; Base has 147 passes. Versions checked: SQLAlchemy 2.0.55, impyla
0.24.0, PyHive 0.7.0, thrift 0.25.0 and thrift-sasl 0.4.3. Staged-file and
branch-file pre-commit plus explicit mypy each exited 0.
Sibling search: `rg -n
'abandon\(|_abandon_cursor\(|_abandon_connection\(|_abandon_service\(|_exception\s*=|is_abandoned|_superset_abandoned|guard_engine_connections\('
superset/db_engine_specs/{impala,hive,base}.py superset/utils/hive.py`, plus
`rg -n '\.abandon\(|ThriftClientGuard' superset --glob '*.py'`.
Siblings checked:
- superset/db_engine_specs/impala.py:189 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:208 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:244 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:410 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:417 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:419 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:449 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/impala.py:451 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/hive.py:229 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/hive.py:256 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/hive.py:272 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/hive.py:276 — fixed in 7ea9cdbe1b
- superset/utils/hive.py:92 — fixed in 7ea9cdbe1b
- superset/utils/hive.py:97 — fixed in 7ea9cdbe1b
- superset/utils/hive.py:101 — fixed in 7ea9cdbe1b
- superset/utils/hive.py:105 — fixed in 7ea9cdbe1b
- superset/utils/hive.py:158 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/base.py:2092 — not affected: the generic
implementation is a no-op and holds no guard or pool state.
- superset/db_engine_specs/base.py:2108 — fixed in 7ea9cdbe1b
- superset/db_engine_specs/base.py:1404 — not affected: this is
exception-class mapping, not a Thrift guard latch.
- superset/models/core.py:816 — not affected: engine creation registers the
driver hooks; per-record invalidation preserves the cached engine and its
configured pool.
- superset/mcp_service/worker.py:626 — not affected: this is WorkerCall
abandonment, not a ThriftClientGuard or DB-API pool-record latch.
Unknown: live warehouse behavior and end-to-end DB_CONNECTION_MUTATOR
configuration were not exercised; these are real SQLAlchemy pools and real
driver transports against local HiveServer2 stand-ins. The newly triggered
GitHub CI run is not claimed complete. I have not resolved the thread or
requested review.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
