aminghadersohi commented on code in PR #45845:
URL: https://github.com/apache/superset/pull/45845#discussion_r4430505171
##########
superset/db_engine_specs/impala.py:
##########
@@ -136,30 +139,43 @@ def execute(
@classmethod
def fetch_data(cls, cursor: Any, limit: int | None = None) ->
list[tuple[Any, ...]]:
"""Wait for asynchronous operations using the public cursor API."""
- if callable(getattr(cursor, "execute_async", None)):
- from impala.error import Error
-
- deadline = time.monotonic() + app.config["SQLLAB_TIMEOUT"]
+ 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:
+ # The operation is released before the time limit propagates so
+ # it does not keep running on the coordinator.
try:
- while cursor.is_executing():
- if time.monotonic() >= deadline:
- cursor.cancel_operation()
- raise TimeoutError("Timed out waiting for the Impala
operation")
- time.sleep(0.1)
- if cursor.execution_failed():
- # Fetching surfaces the driver's detailed asynchronous
error,
- # including errors on statements that produce no result
set.
- cursor.fetchall()
- except Error as ex:
- raise cls.get_dbapi_mapped_exception(ex) from ex
- return super().fetch_data(cursor, limit)
+ cursor.cancel_operation()
+ except TIME_LIMIT_EXCEPTIONS:
+ raise
+ except Exception: # pylint: disable=broad-except
+ logger.warning("cancel_operation() failed after a time limit")
+ raise
Review Comment:
My two earlier replies on this thread were wrong. Their tests mocked
`reconnect()` and only checked call order, so they missed the real problem:
impyla's `HS2Service.reconnect()` only closes and reopens the socket or HTTP
connection, while thrift's `TBufferedTransport` and thrift_sasl's
`TSaslClientTransport` keep their private read and write buffers across
`close()`/`open()`. The unread rest of the interrupted reply was still parsed
as the CancelOperation reply.
Standalone repro with the installed thrift 0.24.0 (real `TSocket` +
`TBufferedTransport`, local server sends `OLD-REPLY-...` on the first
connection and `NEW-REPLY` on the second; read 4 bytes, close, open, read 9):
```
read before interrupt: b'OLD-'
first read after reconnect: b'REPLY-OLD'
```
Fix in ca8cbd0680a8a1d5a3748aa8759737cb34b0b5ee:
`_cancel_operation_after_time_limit` calls a new `_reset_transport`, which
walks the impyla transport stack (`TBufferedTransport` for NOSASL and HTTP,
`TSaslClientTransport` for PLAIN/LDAP/GSSAPI, then `TSocket` or
`ImpalaHttpClient`), checks every buffer attribute exists before changing
anything, closes the transport, replaces every read and write buffer with an
empty one, and reopens it. If the stack is not recognised or the reset raises,
no cancel RPC is sent on it: a warning is logged and the time-limit exception
still propagates. I chose this over a second connection because `fetch_data`
only has the cursor, and over skipping the cancel because the operation would
keep running on the coordinator.
Regression tests, in the new
`tests/unit_tests/db_engine_specs/test_impala_transport.py`. They use impyla's
own `connect()` against a local HiveServer2 stand-in over real binary
(`TBufferedTransport`), SASL PLAIN (`TSaslClientTransport`) and HTTP
(`ImpalaHttpClient`) transports, raise the time limit right after a reply
header is read, and assert exactly one CancelOperation for the right operation
handle, a CloseOperation after it, a new connection for binary and SASL, and no
cancel/close warning:
- `test_handle_cursor_cancels_on_a_clean_transport_after_interrupted_reply`
[binary, sasl, http] x [status, get_log] x [SoftTimeLimitExceeded,
SupersetTimeoutException]
- `test_fetch_data_cancels_on_a_clean_transport_after_interrupted_reply`
[binary, sasl, http] x [status (is_executing poll), rows (FetchResults)] x
[SoftTimeLimitExceeded, SupersetTimeoutException]
- `test_reset_transport_refuses_an_unknown_transport`
Against impala.py at 24860720185c06394e8eb4ccf5c0f4a30cea21b2 all 25 fail:
on binary and SASL the CancelOperation read blocks until the test's 2 s socket
timeout (`TTransportException: read timeout`) and impyla sends it again, and
over HTTP `cancel_operation()` fails. With this commit they pass (25 passed in
5.20s). Mock tests in `test_impala.py` now assert reset-before-cancel
(`test_fetch_data_resets_transport_before_cancelling_an_interrupted_fetch`,
`test_handle_cursor_resets_transport_before_cancelling_an_interrupted_poll`)
and that no cancel RPC is sent when the reset fails
(`test_time_limit_skips_cancel_rpcs_without_a_clean_transport`). test_impala,
test_impala_pending, test_impala_transport, test_hive, test_presto: 273 passed.
The docs paragraph in `configuring-superset.mdx` now says what is actually done.
Siblings checked:
- superset/db_engine_specs/impala.py:163 — fixed in
ca8cbd0680a8a1d5a3748aa8759737cb34b0b5ee (`fetch_data` time-limit handler,
status poll and row fetch)
- superset/db_engine_specs/impala.py:350 — fixed in
ca8cbd0680a8a1d5a3748aa8759737cb34b0b5ee (`handle_cursor` time-limit handler,
status and log polls)
- superset/db_engine_specs/impala.py:150 — not affected: the `fetch_data`
deadline cancel runs after `is_executing()` returned, so no reply is half read
- superset/db_engine_specs/impala.py:319 — not affected: the `handle_cursor`
stop cancel runs after `cursor.status()` returned, so no reply is half read
- superset/sql_lab.py:375 — not affected: the `SoftTimeLimitExceeded`
handler sends no RPC, and closing the impyla connection afterwards only closes
the transport (`HiveServer2Connection.close()`), while the cursor is already
closed so its `__del__` sends nothing
- superset/db_engine_specs/hive.py:465 — not affected: the `fetch_logs`
time-limit branch re-raises without sending an RPC (pyhive's connection close
is covered under Remaining)
- superset/db_engine_specs/hive.py:459 — not affected: the stop cancel runs
after `cursor.poll()` returned, not on a time-limit path
- superset/db_engine_specs/presto.py:1437 — not affected: the stop cancel
runs on an HTTP client with no shared Thrift buffers, not on a time-limit path
Remaining: (1) not run against a real Impala. Untested there: whether the
coordinator accepts CancelOperation for the same session on the reopened binary
connection (impyla's own retries rely on this; if the coordinator closed the
session when the old connection dropped, it cancels the query itself and the
cancel RPC fails with a logged warning), and the GSSAPI/SASL-with-QOP case,
which uses the same `TSaslClientTransport` buffers but was not exercised. (2)
Hive, not changed in this PR and found by reading the code, not by a test:
pyhive's `Connection.close()` sends CloseSession on the same transport when
`get_raw_connection` exits, so a Hive time limit inside a reply (for example
inside `cursor.poll()`, which is the same on master) can read the same stale
bytes.
--
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]
