Repository navigation
fix(impala): end SQL Lab runs that hit a time limit instead of hanging - #45845
aminghadersohi wants to merge 7 commits into
Conversation
A SQL Lab time limit is raised once into whatever the worker is executing: the synchronous SQLLAB_TIMEOUT alarm raises SupersetTimeoutException and the Celery soft limit raises SoftTimeLimitExceeded. Impala's handle_cursor only treated the soft limit specially, so: - a synchronous timeout raised while polling fell into the broad `except Exception` and was swallowed, and the Run kept waiting for the whole query before failing from fetch_data's own deadline; - either limit raised while the progress log was fetched was swallowed as a failed GetLog() call, and with the limit already spent nothing stopped the query until the hard limit; - a limit raised while fetch_data waited on the operation left it running on the coordinator. Both limits are now cancelled and propagated from handle_cursor, including from the progress log call and the cancel helper, and fetch_data cancels the operation before propagating them. Hive's handle_cursor had the same swallow around fetch_logs and now lets the limits through. Presto and Trino do not wrap their polling in a broad except.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #45845 +/- ##
===========================================
+ Coverage 65.38% 84.71% +19.33%
===========================================
Files 3069 3069
Lines 205465 205770 +305
Branches 48310 48358 +48
===========================================
+ Hits 134336 174320 +39984
+ Misses 68043 28152 -39891
- Partials 3086 3298 +212
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Product acceptance at Works as described for the phases the PR covers (a limit that lands while the operation is running or admission-queued):
Gap in Unchanged and pre-existing, as the description says: a threaded sync request (gthread) still has no time limit (no response at 45 s, here and at the merge-base). Not covered: the Hive change was not run against a live Hive server. Testing note: Impala's |
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
|
rebenitez1802
left a comment
rebenitez1802
left a comment
There was a problem hiding this comment.
Request changes: the fix is right and well tested (16 of 17 new tests fail on master; on real Impala 4.5.2 a sync timeout now ends the run at the limit instead of about 20 s later), but the new cancel in the row-fetch time-limit handler can hang the worker. Details and a live-validated fix are inline.
| # 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 |
There was a problem hiding this 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 staysrunning. - 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:
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()| # 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 | |
| # 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 |
There was a problem hiding this comment.
I addressed this in bde3fa4: the time-limit handler reconnects before cancellation and uses _cancel_operation to release the operation and cursor even when cancellation fails.
Regression test: tests/unit_tests/db_engine_specs/test_impala.py::test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch (cancel success and failure). Before the change, both cases failed at assert calls.mock_calls[:2] == [call.reconnect(), call.cancel_operation()]: At index 0 diff: call.cancel_operation() != call.reconnect(). After the change, the Impala, Impala pending-operation, and Hive unit suites passed: 93 passed. Pre-commit passed, including mypy. Real Impala not run; I have not independently remeasured the live hang or coordinator resource release.
Siblings checked:
- superset/db_engine_specs/impala.py:140 — fixed in bde3fa4
- superset/db_engine_specs/impala.py:178 — not affected: shared cleanup releases handles after an ordinary cancel failure; the interrupted-fetch caller reconnects before entering it.
- superset/db_engine_specs/impala.py:203 — not affected: polling and explicit Stop occur before row fetching, not after an interrupted FetchResults response.
- superset/db_engine_specs/impala.py:321 — not affected: cancel_query uses a separate HTTP request, not the cursor transport.
- superset/db_engine_specs/base.py:1624 — not affected: generic fetching propagates exceptions without issuing a cancellation RPC.
- superset/db_engine_specs/hive.py:200 — not affected: Hive fetching does not cancel after a time limit; its log-poll handler rethrows the limit without a cancellation RPC.
- superset/db_engine_specs/spark.py — not affected: no fetch override or cancel_operation call.
- superset/sql_lab.py:369 — not affected: dispatches fetching through the engine spec and handles timeouts without a cursor cancellation RPC.
There was a problem hiding this comment.
Fixed the remaining polling path in 2486072. My earlier reply wrongly marked handle_cursor (then impala.py:203) unaffected. A signal can interrupt any RPC reply, including GetOperationStatus and GetLog, not just FetchResults. Both time-limit handlers use one helper to reconnect, cancel, and release operation/cursor handles before re-raising; non-time-limit paths are unchanged.
Regression test: test_handle_cursor_reconnects_before_cancelling_an_interrupted_poll exercises status/get_log, SoftTimeLimitExceeded/SupersetTimeoutException, and successful/failed cancellation. Before the code change, PYTHONPATH=.:superset-core/src pytest -q tests/unit_tests/db_engine_specs/test_impala.py -k interrupted_poll returned 8 failed, 29 deselected; the assertion was At index 0 diff: call.cancel_operation() != call.reconnect(). After the change, the Impala/Hive suite returned 101 passed, including that test and test_fetch_data_reconnects_before_cancelling_an_interrupted_fetch. Base/Presto neighbors returned 282 passed. Staged-file and branch-changed-file pre-commit checks passed, including mypy. Real Impala not run; no live transport or server-side slot-release claim.
Search: rg -n 'cancel_operation|_cancel_operation|SoftTimeLimitExceeded|TIME_LIMIT_EXCEPTIONS|def handle_cursor|def fetch_data' superset/db_engine_specs superset/sql_lab.py, plus Spark inheritance.
Siblings checked:
- superset/db_engine_specs/impala.py:140 — fixed in 2486072; fetch_data uses the shared reconnect-before-cancel helper for time limits during polling or row fetching.
- superset/db_engine_specs/impala.py:208 — fixed in 2486072; handle_cursor routes interrupted status and log RPCs through the same helper.
- superset/db_engine_specs/impala.py:167 — fixed in 2486072; the shared helper reconnects before _cancel_operation, preserving reconnect-failure logging and propagation of another time limit.
- superset/db_engine_specs/impala.py:183 — not affected: the low-level cancel/release helper is reached through reconnect for signal-driven cancellation; ordinary RPC failures are logged, and another time limit propagates rather than starting another RPC.
- superset/db_engine_specs/impala.py:150 — not affected: the monotonic deadline is checked after is_executing returns, with no interrupted reply; signal exceptions instead enter the reconnect handler.
- superset/db_engine_specs/impala.py:237 — not affected: explicit Stop is checked after status returns, not from a signal interrupt; a time limit during cancellation enters the outer reconnect handler.
- superset/db_engine_specs/impala.py:323 — not affected: cancel_query uses a separate HTTP requester, not the interrupted impyla transport.
- superset/db_engine_specs/base.py:144 — not affected: TIME_LIMIT_EXCEPTIONS only defines the exception tuple; it sends no RPC.
- superset/db_engine_specs/base.py:1624 — not affected: generic fetching maps exceptions but issues no cancellation RPC after interruption; Impala wraps it with the shared handler.
- superset/db_engine_specs/base.py:1670 — not affected: pagination entry point raises NotImplementedError and sends no RPC.
- superset/db_engine_specs/base.py:2067 — not affected: default handle_cursor is a no-op.
- superset/db_engine_specs/hive.py:200 — not affected: PyHive fetching has no signal-driven cancel follow-up on the interrupted transport.
- superset/db_engine_specs/hive.py:436 — not affected: explicit Stop calls cancel after poll returns; the TIME_LIMIT_EXCEPTIONS handler at line 465 re-raises without a cancellation RPC.
- superset/db_engine_specs/spark.py:51 — not affected: inherits Hive handling, which re-raises time limits without following them with cancellation on the interrupted transport.
- superset/sql_lab.py:375 — not affected: marks the query timed out and raises an application error; it does not send a cancellation RPC on the interrupted transport.
- superset/db_engine_specs/crate.py:99 — not affected: post-processes generic fetch results without a time-limit cancel handler.
- superset/db_engine_specs/exasol.py:165 — not affected: converts generic fetch rows without a time-limit cancel handler.
- superset/db_engine_specs/elasticsearch.py:274 — not affected: HTTP cursor pagination, not impyla signal-driven cancellation.
- superset/db_engine_specs/elasticsearch.py:392 — not affected: OpenDistro HTTP cursor pagination, not impyla signal-driven cancellation.
- superset/db_engine_specs/mssql.py:194 — not affected: converts ODBC rows without a time-limit cancel handler.
- superset/db_engine_specs/bigquery.py:416 — not affected: progressive fetching has no cancellation RPC after an interrupted reply.
- superset/db_engine_specs/datastore.py:319 — not affected: converts fetched rows without a time-limit cancel handler.
- superset/db_engine_specs/ocient.py:322 — not affected: fetch failure removes an in-memory mapping and re-raises, without cancellation.
- superset/db_engine_specs/ocient.py:383 — not affected: records the query identifier then delegates to the no-op base handle_cursor.
- superset/db_engine_specs/drill.py:187 — not affected: handles an empty-result RuntimeError without issuing cancellation.
- superset/db_engine_specs/trino.py:331 — not affected: Trino cursor handling does not use impyla or cancel in a signal-exception handler.
- superset/db_engine_specs/presto.py:1415 — not affected: explicit Stop cancellation follows a completed poll, not an interrupted reply.
SUMMARY
An Impala query that hits a SQL Lab time limit could keep running instead of ending as a failed, timed-out run.
SQL Lab raises its time limits into whatever the worker is executing: the main-thread synchronous
SQLLAB_TIMEOUTalarm raisesSupersetTimeoutException, and the Celery soft limit raisesSoftTimeLimitExceeded. Each fires once.ImpalaEngineSpec.handle_cursoronly handled the soft limit around its pollingsleep/statuscalls, so:except Exception: returnand was swallowed. The Run then waited for the whole query and failed fromfetch_data's own deadline instead of atSQLLAB_TIMEOUT.GetLog()call. The limit had already fired, so nothing stopped the query until the Celery hard limit, and the row stayedrunninguntil then.fetch_datawaited on the operation propagated without cancelling it, leaving it running on the coordinator and holding its admission slot.Changes:
TIME_LIMIT_EXCEPTIONS(the two exceptions above) is defined once indb_engine_specs/base.py.ImpalaEngineSpec.handle_cursorcancels the operation and re-raises on either limit, and no longer swallows them in theget_log()call._cancel_operationlets them through its cleanup RPCs.ImpalaEngineSpec.fetch_datacancels the operation before re-raising either limit.HiveEngineSpec.handle_cursorhad the same swallow aroundfetch_logs()and now lets the limits through.Checked and unchanged: Presto and Trino do not wrap their polling in a broad
except, and the base and Ocienthandle_cursorhave no polling loop.Not changed here:
get_sql_resultsstill hardcodessoft_time_limit=21600/time_limit=21660in the@celery_app.taskdecorator, soSQLLAB_ASYNC_TIME_LIMIT_SEConly affects the error message. Applying it to the task limits changes behaviour for every engine and needs a hook at app initialization, so it is left for a separate change. Operators can already settask_annotationsforsql_lab.get_sql_resultsinCELERY_CONFIG.BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
N/A
TESTING INSTRUCTIONS
Unit tests with a mocked Impala/Hive cursor; no live Impala check was run:
At commit
85ef2d19eb32be18d0531806302f5312320ce037,PYTHONPATH=.:superset-core/src python -m pytest tests/unit_tests/db_engine_specs/test_impala.py tests/unit_tests/db_engine_specs/test_impala_pending.py tests/unit_tests/db_engine_specs/test_hive.pypassed: 86 tests. Removing only the nestedexcept TIME_LIMIT_EXCEPTIONS: raisefromfetch_datamade both cases oftest_fetch_data_propagates_time_limit_during_cancelfail; restoring it returned the full suite to 86 passing tests. The new test asserts the cancellation exception's identity, distinguishing it from the original wait exception.test_status_failure_still_ends_polling_without_raisingpreserves ordinary RPC-error behavior; it is not a timeout regression test.test_sync_timeout_fails_the_run_and_cancels_the_operationdrivesexecute_queryand asserts propagation asSqlLabExceptionwith the timeout message and operation cancellation; it does not assert a persisted database status.The synchronous alarm only runs on the main thread (
superset/utils/core.py:835-839). Threaded synchronous requests, including the default gthread server, have no polling deadline inhandle_cursor; this pre-existing gap is not covered here. Windows timer timeouts also remain outside this change: they raiseKeyboardInterruptinside the timed block before conversion on context exit.Manual check against Impala (main-thread synchronous execution only): run a long query in SQL Lab with a small
SQLLAB_TIMEOUT(sync) or a Celery soft time limit (async) and confirm the query fails at the limit and its operation is no longer running in the Impala web UI.ADDITIONAL INFORMATION