aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4278509073
##########
superset/sql/execution/celery_task.py:
##########
@@ -481,25 +570,60 @@ def _execute_sql_statements(
)
) from ex
- # Check if stopped
- if not execution_results:
- payload.update({"status": QueryStatus.STOPPED.value})
- return payload
+ # Cancelling can clear cursor.description, leaving a nonempty result
+ # list that looks like successful DML. Check intent, not just results.
+ stopped = not execution_results or check_stopped()
# Commit for mutations
- if parsed_script.has_mutation() or query.select_as_cta:
+ committed = False
+ if not stopped and (parsed_script.has_mutation() or
query.select_as_cta):
Review Comment:
Thanks, this was valid at 24bca1e910010632f50dc19c477fca0ef94625c3. Fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7.
The commit decision now asks the connection whether it commits each
statement: `BaseEngineSpec.is_autocommit(conn)`
(superset/db_engine_specs/base.py:3028) reads `autocommit` (psycopg2, pyodbc),
`get_autocommit()` (MySQL drivers) or an `AUTOCOMMIT` isolation level (Trino).
Drivers that report none of these use `autocommit_by_default`, which is True
for the PyHive specs (Presto, Hive and its descendants) and for impyla
(Impala), whose `commit()` does nothing. On such a connection, a mutation or
CTAS that finished executing counts as committed work: `on_commit` fires and
the query ends SUCCESS even if a stop was recorded after it finished.
A stop still wins if it was loaded into the worker's session while the
statements ran. Either the check between statements loads it, or the engine's
polling does: Impala, Hive and Presto `handle_cursor` load it, then cancel the
running statement. `stop_seen_during_execution`
(superset/sql/execution/celery_task.py:573) looks only at attributes already
loaded and does not read the row again. That is what separates a stop that
interrupted execution from one that arrived after the statement finished.
Regression tests (each [autocommit-after_finish-*] case fails at
24bca1e910010632f50dc19c477fca0ef94625c3 with `assert 'stopped' ==
<QueryStatus.SUCCESS: 'success'>` and passes at
0eab948613546ed4a7a678c7b19c5fcdc3220ab7). The old-head run used 24bca1e9's
`celery_task.py` and `sql_lab.py`, with the new tests and the engine-spec
detection copied in:
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_autocommit_work_finished_before_a_stop_is_committed`:
16 cases (autocommit or transaction × stop after the statement finished or
loaded while it ran × STOPPED or early-cancel intent × mutation or CTAS), with
the stop committed from a second SQLite session. The 4
[autocommit-after_finish-*] cases are the regression; the other 12 pass on both
heads and end STOPPED without `conn.commit()`.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_autocommit_work_finished_before_a_stop_is_committed`:
the same for the legacy path. [autocommit-after_finish-mutation] and
[autocommit-after_finish-ctas] are the regression.
- `tests/unit_tests/db_engine_specs/test_base.py::test_is_autocommit`,
`test_is_autocommit_real_connections`,
`test_is_autocommit_falls_back_to_the_engine_default`, and
`tests/unit_tests/db_engine_specs/test_trino.py::test_is_autocommit_reads_trino_isolation_level`
(a real `trino.dbapi.Connection`) cover the detection.
Siblings checked:
- superset/sql_lab.py:762 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 (the legacy target-commit check uses
the same rule)
- superset/sql/execution/celery_task.py:712 — not affected: transactional
connections keep the stop check before `conn.commit()`
- superset/sql/execution/celery_task.py:745 — not affected:
`_publish_query_result(durable=True)` already publishes committed work over a
later stop
- superset/sql/execution/executor.py:787 — not affected: SQLExecutor's sync
path passes no stop check and never calls `conn.commit()`
- superset/db_engine_specs/presto.py:171 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 (PyHive reports no transaction mode,
so the spec default applies; Trino is read from the connection)
- superset/db_engine_specs/impala.py:54 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 (impyla reports no transaction mode)
Remaining:
- In a multi-statement script on an autocommit engine, statements that
finished before the stop are durable, but the query ends STOPPED because the
stop interrupted the script. The status does not say which statements applied.
- If the engine finishes the statement between its last status poll and the
cancel call, the cancel does nothing, but the query ends STOPPED with the rows
written. The window is the time between `cursor.status()` and the cancel RPC in
`handle_cursor`.
- On the legacy path, `execute_query` reloads the row just before it
executes. A stop committed after the per-block check but before that reload is
loaded too, so on an engine that does not poll for stops the statement still
runs and the query ends STOPPED with its rows written.
- Drivers that report no transaction mode on engines whose spec does not set
`autocommit_by_default` (for example BigQuery, or Snowflake, whose connector
exposes `autocommit` as a method) behave as before: a stop before the commit
check ends STOPPED with the work already applied.
##########
superset/sql_lab.py:
##########
@@ -825,7 +867,9 @@ def execute_sql_statements( # noqa: C901
# unit-test suite driving this fix never exercised.
db.session.flush()
db.session.refresh(query)
- if query.status not in (QueryStatus.FAILED, QueryStatus.STOPPED):
+ if query.status != QueryStatus.FAILED and (
+ committed or query.status != QueryStatus.STOPPED
+ ):
query.status = QueryStatus.SUCCESS
Review Comment:
Thanks, this was valid at 24bca1e910010632f50dc19c477fca0ef94625c3: the row
stayed STOPPED, but the returned payload still said SUCCESS and carried the
data, and the STOPPED row referenced the stored results. Fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7.
The final check in `execute_sql_statements` (superset/sql_lab.py:971)
reloads only the status, without flushing. If a stop landed while the results
were serialized or written to the results backend, and the target database did
not commit, it now:
- reloads the row so the pending success fields are dropped and the stop's
end time is kept;
- deletes the stored results, because the stopped row will not reference
them;
- commits, and returns `{"query_id": ..., "status": "stopped"}`. That is the
same minimal payload as the write-failure stop return at :931.
Regression test:
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_stop_during_select_finalization_ends_stopped`,
5 cases. A second SQLite session commits the stop during
`_serialize_and_expand_data` (no store; store and return; store only) or during
`results_backend.set` (store and return; store only). All 5 fail at
24bca1e910010632f50dc19c477fca0ef94625c3 (old `sql_lab.py`, new tests). For
example, [serialize-False-True] fails with `assert {'columns': [...],
...SUCCESS...} == {'query_id': ..., 'status': 'stopped'}` and
[backend_write-True-False] with `assert None == {'query_id': 123, 'status':
<QueryStatus.STOPPED: 'stopped'>}`. All 5 pass at
0eab948613546ed4a7a678c7b19c5fcdc3220ab7. They also check that the row is
STOPPED with the stop's `end_time`, `results_key` is None, and the stored blob
was deleted.
Changed test:
`test_get_sql_results_keeps_published_results_of_stopped_query` relied on a
stopped row keeping its key. It is now
`test_get_sql_results_keeps_published_results_after_final_commit`, without a
stop: once the SUCCESS row with the key is committed, a later failure ends
FAILED and keeps the key and the blob.
Siblings checked:
- superset/sql/execution/celery_task.py:759 — not affected: the Celery
publish check returns the minimal STOPPED payload and deletes the blob
- superset/sql/execution/celery_task.py:765 — not affected: when the guarded
publish loses to a stop, it does the same
- superset/sql_lab.py:810 — not affected: the check before the SUCCESS
payload is built already returns STOPPED
- superset/sql_lab.py:919 — not affected: the write-failure branch already
returns the minimal STOPPED payload
- superset/sql_lab.py:730 — not affected: a stop raised by the statement
returns STOPPED before any payload is built
Remaining: a stop committed between this status read and the final commit is
still overwritten with SUCCESS. The payload and the row then agree on SUCCESS.
The legacy path has no guarded UPDATE here (the Celery path does), and this
window is disclosed in the code comment.
##########
superset/sql/execution/celery_task.py:
##########
@@ -198,102 +218,110 @@ def _finalize_successful_query(
}
)
- query.rows = total_rows
- query.progress = 100
- query.set_extra_json_key("progress", None)
- # Store columns from last statement (for compatibility)
- if execution_results and execution_results[-1][1] is not None:
- query.set_extra_json_key("columns", execution_results[-1][1].columns)
- query.end_time = now_as_float()
-
- payload.update(
- {
- "status": QueryStatus.SUCCESS.value,
- "statements": statements_data,
- "total_execution_time_ms": total_execution_time_ms,
- "query": query.to_dict(),
- }
- )
- payload["query"]["state"] = QueryStatus.SUCCESS.value
+ with db.session.no_autoflush:
+ query.rows = total_rows
+ query.progress = 100
+ query.set_extra_json_key("progress", None)
+ # Store columns from last statement (for compatibility)
+ if execution_results and execution_results[-1][1] is not None:
+ query.set_extra_json_key("columns",
execution_results[-1][1].columns)
+ query.end_time = now_as_float()
+
+ payload.update(
+ {
+ "status": QueryStatus.SUCCESS.value,
+ "statements": statements_data,
+ "total_execution_time_ms": total_execution_time_ms,
+ "query": query.to_dict(),
+ }
+ )
+ payload["query"]["state"] = QueryStatus.SUCCESS.value
def _store_results_in_backend(
query: Query,
payload: dict[str, Any],
database: Any,
+ durable: bool = False,
) -> None:
"""Store query results in the results backend."""
- key = str(uuid.uuid4())
- payload["query"]["resultsKey"] = key
- logger.info(
- "Query %s: Storing results in results backend, key: %s",
- str(query.id),
- key,
- )
- stats_logger = app.config["STATS_LOGGER"]
- with stats_timing("sqllab.query.results_backend_write", stats_logger):
- with stats_timing(
- "sqllab.query.results_backend_write_serialization", stats_logger
- ):
- serialized_payload = _serialize_payload(payload)
-
- # Check payload size limit
- if sql_lab_payload_max_mb :=
app.config.get("SQLLAB_PAYLOAD_MAX_MB"):
- serialized_payload_size = len(serialized_payload)
- max_bytes = sql_lab_payload_max_mb * BYTES_IN_MB
-
- if serialized_payload_size > max_bytes:
- logger.info("Result size exceeds the allowed limit.")
- raise SupersetErrorException(
- SupersetError(
- message=(
- f"Result size "
- f"({serialized_payload_size / BYTES_IN_MB:.2f}
MB) "
- f"exceeds the allowed limit of "
- f"{sql_lab_payload_max_mb} MB."
- ),
-
error_type=SupersetErrorType.RESULT_TOO_LARGE_ERROR,
- level=ErrorLevel.ERROR,
+ with db.session.no_autoflush:
+ key = str(uuid.uuid4())
+ payload["query"]["resultsKey"] = key
+ logger.info(
+ "Query %s: Storing results in results backend, key: %s",
+ str(query.id),
+ key,
+ )
+ stats_logger = app.config["STATS_LOGGER"]
+ with stats_timing("sqllab.query.results_backend_write", stats_logger):
+ with stats_timing(
+ "sqllab.query.results_backend_write_serialization",
stats_logger
+ ):
+ serialized_payload = _serialize_payload(payload)
+
+ # Check payload size limit
+ if sql_lab_payload_max_mb :=
app.config.get("SQLLAB_PAYLOAD_MAX_MB"):
+ serialized_payload_size = len(serialized_payload)
+ max_bytes = sql_lab_payload_max_mb * BYTES_IN_MB
+
+ if serialized_payload_size > max_bytes:
+ logger.info("Result size exceeds the allowed limit.")
+ raise SupersetErrorException(
+ SupersetError(
+ message=(
+ f"Result size "
+ f"({serialized_payload_size /
BYTES_IN_MB:.2f} MB) "
+ f"exceeds the allowed limit of "
+ f"{sql_lab_payload_max_mb} MB."
+ ),
+
error_type=SupersetErrorType.RESULT_TOO_LARGE_ERROR,
+ level=ErrorLevel.ERROR,
+ )
)
- )
-
- cache_timeout = database.cache_timeout
- if cache_timeout is None:
- cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"]
- compressed = zlib_compress(serialized_payload)
- logger.debug("*** serialized payload size: %i",
len(serialized_payload))
- logger.debug("*** compressed payload size: %i", len(compressed))
-
- write_success = results_backend.set(key, compressed, cache_timeout)
- if not write_success:
- logger.error(
- "Query %s: Failed to store results in backend, key: %s",
- str(query.id),
- key,
- )
- stats_logger.incr("sqllab.results_backend.write_failure")
- query.results_key = None
- query.status = QueryStatus.FAILED
- query.error_message = (
- "Failed to store query results in the results backend. "
- "Please try again or contact your administrator."
- )
- db.session.commit() # pylint: disable=consider-using-transaction
- raise SupersetErrorException(
- SupersetError(
- message=__("Failed to store query results. Please try
again."),
- error_type=SupersetErrorType.RESULTS_BACKEND_ERROR,
- level=ErrorLevel.ERROR,
+ cache_timeout = database.cache_timeout
+ if cache_timeout is None:
+ cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"]
+
+ compressed = zlib_compress(serialized_payload)
+ logger.debug("*** serialized payload size: %i",
len(serialized_payload))
+ logger.debug("*** compressed payload size: %i", len(compressed))
+
+ write_success = results_backend.set(key, compressed, cache_timeout)
+ if not durable and _make_check_stopped_fn(query,
preserve_pending=True)():
Review Comment:
Follow-up: an independent check of my last reply found a case I had only
listed as remaining. A metadata commit can reach the database and still raise,
for example when its acknowledgement is lost. The final cleanup in
`get_sql_results` (both `return_results` modes) and in `execute_sql_task` then
deleted the stored blob, although the committed FAILED or SUCCESS row pointed
at it.
Fixed in 0eab948613546ed4a7a678c7b19c5fcdc3220ab7. Both cleanups
(superset/sql_lab.py:365, superset/sql/execution/celery_task.py:451) now call
`_discard_unreferenced_results` (superset/sql_lab.py:226,
superset/sql/execution/celery_task.py:511). It rolls back, reads the row's
`results_key` in a new transaction, and deletes the blob only if the row does
not reference it. If that read also fails because the metadata database is
unreachable, it keeps the blob and logs "Kept results that may be referenced";
the blob expires with `cache_timeout`.
This changes one case on purpose. When the metadata database stays down, the
blob used to be deleted; now it is kept, because the worker cannot tell whether
a commit applied. An orphaned blob that expires is safer than a row that points
at deleted results.
Regression tests. The old-head runs used
24bca1e910010632f50dc19c477fca0ef94625c3's `celery_task.py` and `sql_lab.py`
with the new tests; each case below fails there with "Expected 'delete' to not
have been called. Called 1 times." and passes at
0eab948613546ed4a7a678c7b19c5fcdc3220ab7:
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_commit_that_applies_and_raises_keeps_referenced_results[failed_row]`:
the error handler's FAILED commit reaches the database and then raises.
- `[published_row]` of the same test: the publish commit reaches the
database and raises, and the metadata database then stays down.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_commit_that_applies_and_raises_keeps_results[failed_row-False]`,
`[failed_row-True]` and `[published_row-True]`: the same on the legacy path. A
second session reads the committed row with the key, and the blob is still
stored.
- `[no_reread]` of
`tests/unit_tests/sql/execution/test_celery_task.py::test_unrecorded_error_after_results_write_deletes_only_unreferenced_results`
and of
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_unrecorded_error_after_write_deletes_unreferenced_results`:
every metadata statement fails, so the blob is kept and the warning logged.
- Controls that pass on both heads: the `[reread-*]` cases of those two
tests (the read works and shows no reference, so the blob is deleted, including
when the delete itself fails), and
`test_error_path_deletes_results_only_without_committed_failure` and
`test_get_sql_results_deletes_results_only_without_committed_failure` (a commit
that did not apply, so the blob is deleted).
- `test_discard_unreferenced_results_without_stored_results` covers the
early return when nothing was stored.
Siblings checked:
- superset/sql/execution/celery_task.py:451 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7
- superset/sql_lab.py:365 — fixed in 0eab948613546ed4a7a678c7b19c5fcdc3220ab7
- superset/sql/execution/celery_task.py:106 — not affected: the stop check
found STOPPED before any commit carried the key
- superset/sql/execution/celery_task.py:137 — not affected: the guarded
UPDATE found STOPPED, so nothing was committed
- superset/sql/execution/celery_task.py:315 — not affected: deleted before
any commit carried the key
- superset/sql/execution/celery_task.py:760 — not affected: deleted before
the publish commit
- superset/sql/execution/celery_task.py:766 — not affected: the guarded
publish committed nothing
- superset/sql_lab.py:151 — not affected: the legacy handler's stop return
commits nothing
- superset/sql_lab.py:976 — not affected: new in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7; the reload drops the pending key
before the delete, so the STOPPED commit does not carry it
Remaining:
- If the metadata database cannot be read after the failure, the blob is
kept until `cache_timeout`, even when no row references it.
- If the worker dies between the write and the publish, or a delete fails,
the blob also stays until `cache_timeout`, as before.
- The legacy `handle_query_error` still lets a recorded stop replace a
timeout and records FAILED rather than TIMED_OUT. That is master's behaviour
(see its TODO), and this PR does not change it.
##########
superset/sql_lab.py:
##########
@@ -628,8 +651,12 @@ def execute_sql_statements( # noqa: C901
return payload
# Commit the connection so CTA queries will create the table and any
DML.
+ committed = False
if parsed_script.has_mutation() or query.select_as_cta:
Review Comment:
Follow-up: my last reply listed autocommit engines as remaining. An
independent check also showed that on connections that commit each statement,
both the legacy and the Celery paths could end STOPPED with the rows or the
table already written, whenever the stop arrived after the statement finished
but before the target-commit check. Fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7.
Both paths now ask the connection whether each statement commits when it
finishes: `BaseEngineSpec.is_autocommit(conn)`, at superset/sql_lab.py:762 and
superset/sql/execution/celery_task.py:702. It reads the driver's `autocommit`,
`get_autocommit()` or an `AUTOCOMMIT` isolation level (Trino). For drivers that
report nothing it falls back to `autocommit_by_default`, which is True for
PyHive (Presto, Hive) and impyla (Impala). On such a connection, finished work
counts as committed and the query ends SUCCESS. A stop ends STOPPED only if it
was loaded into the worker's session while the statements ran: by the check
between statements, or by engine polling that then cancels the running
statement. `stop_seen_during_execution` at
superset/sql/execution/celery_task.py:573 does not read the row again.
Transactional connections keep the check you asked for: a stop before
`conn.commit()` skips it and ends STOPPED.
Regression tests (each [autocommit-after_finish-*] case fails at
24bca1e910010632f50dc19c477fca0ef94625c3 with `assert 'stopped' ==
<QueryStatus.SUCCESS: 'success'>` and passes at
0eab948613546ed4a7a678c7b19c5fcdc3220ab7):
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_autocommit_work_finished_before_a_stop_is_committed`:
mutation and CTAS, with a stop committed by a second SQLite session as the
last statement finishes. The [autocommit-while_running-*] and [transaction-*]
cases pass on both heads and end STOPPED without `conn.commit()`.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_autocommit_work_finished_before_a_stop_is_committed`:
the same for Celery, with both STOPPED and early-cancel intent (16 cases).
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_stop_before_target_commit_stays_stopped`
and `test_get_sql_results_keeps_committed_work_when_stop_races_commit` still
pass.
- `tests/unit_tests/db_engine_specs/test_base.py::test_is_autocommit*` and
`tests/unit_tests/db_engine_specs/test_trino.py::test_is_autocommit_reads_trino_isolation_level`
cover the detection.
Siblings checked:
- superset/sql_lab.py:762 — fixed in 0eab948613546ed4a7a678c7b19c5fcdc3220ab7
- superset/sql/execution/celery_task.py:702 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7
- superset/sql_lab.py:810 — not affected: once `committed` is set, the later
stop checks treat the work as committed, as they already did after a real commit
- superset/sql_lab.py:971 — not affected: the same `committed` flag applies
at the final status check
- superset/sql/execution/executor.py:787 — not affected: SQLExecutor's sync
path has no stop check and never calls `conn.commit()`
Remaining:
- On an autocommit engine, statements in a multi-statement script that
finished before a stop are durable, but the query ends STOPPED. The status does
not say which statements applied.
- If the engine finishes a statement between its last status poll and the
cancel call, the query ends STOPPED with the rows written.
- On the legacy path, a stop committed after the per-block check but before
`execute_query` reloads the row still runs the statement on engines that do not
poll for stops, and ends STOPPED with the rows written.
- Drivers that report no transaction mode on engines without
`autocommit_by_default` (for example BigQuery, Snowflake) behave as before.
--
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]
