rebenitez1802 commented on code in PR #45845:
URL: https://github.com/apache/superset/pull/45845#discussion_r4428289067
##########
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:
🔴 High: time-limit cancel during row fetch reuses a half-read connection and
can hang
The time limit is raised by a signal and usually lands inside impyla's
`FetchResults` read, often partway through the reply.
`cursor.cancel_operation()` then runs over the same connection, and its reply
is read from the leftover row bytes. impyla doesn't check which call a reply
belongs to, so this produces either a garbage-sized read (0.5 to 1.9 GB
requested) or `TypeError: Error while skipping unknown field`. There is no
socket timeout by default.
On real Impala 4.5.2 with a ~200 MB reply, 7 of 8 attempts hung:
- Once inside `fetch_data`, so the worker sits until the Celery hard limit
or the gunicorn timeout, and the query row stays `running`.
- Six times in the cursor's `__del__` after the cancel had failed, which
also leaves the query running.
On master, 0 of 7 hung, because no RPC follows the limit
(`get_raw_connection`'s `closing()` only closes the transport).
Reopening the connection before cancelling fixed it: 8 of 8 clean live runs,
the other time-limit scenarios unchanged, and the fetch-phase timeout dropping
from about 6 s to 0.8 s. It also routes this handler through
`_cancel_operation`, so the operation and cursor are closed when the cancel
fails. Unit tests: 92 passed, with one test added:
```python
def test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch(
app_context: None, impala_errors: None
) -> None:
"""A limit inside FetchResults leaves its reply unread; cancel on a new
socket."""
from unittest.mock import MagicMock
from superset.db_engine_specs.base import BaseEngineSpec
calls = MagicMock()
cursor = Mock()
cursor.is_executing.return_value = False
cursor.execution_failed.return_value = False
cursor.session.service.reconnect = calls.reconnect
cursor.cancel_operation = calls.cancel_operation
with patch.object(
BaseEngineSpec, "fetch_data", side_effect=SoftTimeLimitExceeded()
):
with pytest.raises(SoftTimeLimitExceeded):
spec.fetch_data(cursor, 10)
assert [c[0] for c in calls.mock_calls[:2]] == ["reconnect",
"cancel_operation"]
cursor.close_operation.assert_called_once_with()
cursor.close.assert_called_once_with()
```
```suggestion
# The operation is released before the time limit propagates so
# it does not keep running on the coordinator. The limit can land
# part way through a FetchResults response, so the transport is
# reopened first: reading the cancel's reply from the remainder
of
# that response would hang or request a garbage-sized read.
try:
cursor.session.service.reconnect()
except TIME_LIMIT_EXCEPTIONS:
raise
except Exception: # pylint: disable=broad-except
logger.warning("reconnect() failed after a time limit")
cls._cancel_operation(cursor)
raise
```
--
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]
