aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4277968643
##########
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: my last reply was wrong about one window, and its claim that a
no-stop failure never strands the blob was wrong too. Celery recorded the key
only after the status read that follows `results_backend.set`, so if that read
raised, the key was lost: with a concurrent stop the query ended STOPPED and
the blob stayed behind, and without a stop the FAILED row had
`results_key=None` and nothing referenced the blob.
Fixed in 88a85e28e8047bc2d3029e1e786285360f8160e7:
- `_store_results_in_backend` calls `on_results_key(key)` right after a
successful write, before any metadata read, and `execute_sql_task` passes the
tracked key to `_handle_query_error`. The handler deletes the blob when a stop
wins, and otherwise writes the key back onto the FAILED row, so a rollback
cannot drop it.
- A flush failure in the publish commit left the session unusable, so the
error lookup itself raised and the task crashed with the row still RUNNING.
`_get_query_after_error` rolls the session back and retries the lookup once.
- If the failure still cannot be recorded (metadata database down),
`execute_sql_task` deletes the tracked blob and re-raises.
- Legacy: `handle_query_error` writes the tracked key back onto the FAILED
row, and `get_sql_results` deletes the tracked blob if the lookup or the
handler fails.
Regression tests (each fails at e7ebebc34b8c0281625d79b83d9ac5c072470e6c and
passes at 88a85e28e8047bc2d3029e1e786285360f8160e7):
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_error_after_results_write_deletes_stored_results_when_stopped[status_read-*]`:
the first SELECT after the write raises. The 4 stop cases failed with
"Expected 'delete' to be called once. Called 0 times."; the 2 no-stop cases
failed with `assert None == '<key>'` on `results_key`. Now the stop cases end
STOPPED with the blob deleted, and the no-stop cases end FAILED with the row
keeping the key.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_error_after_results_write_deletes_stored_results_when_stopped[publish_flush-*]`:
the ORM flush in the publish commit raises. All 4 failed with
`PendingRollbackError`.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_unrecorded_error_after_results_write_deletes_stored_results`:
every metadata statement after the write raises. Both cases failed with 0
deletes.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_error_after_write_deletes_results_when_stopped[final_flush-no_stop-*]`:
the legacy final flush raises and the session rolls back. All 4 failed with
`assert None == '<key>'`. The 4 stop cases with the same failure already
deleted the blob at e7ebebc34b8c0281625d79b83d9ac5c072470e6c.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_unrecorded_error_after_write_deletes_results`:
the legacy lookup gives up after its retries. Both cases failed with 0 deletes.
- The delete-failure warning reads "Failed to delete unpublished results",
because it also covers unrecorded failures; the existing assertions were
updated to match.
- `tests/unit_tests/sql/execution/test_celery_task.py`,
`tests/unit_tests/sql_lab_test.py` and
`tests/unit_tests/test_sqllab_error_stacktrace.py`: 191 passed, 2 skipped.
Siblings checked:
- superset/sql/execution/celery_task.py:310 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (the key is tracked right after `set`
succeeds)
- superset/sql/execution/celery_task.py:311 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (if the post-write status read raises,
the handler gets the tracked key, see the status_read tests; if it finds a
stop, :313 deletes the blob)
- superset/sql/execution/celery_task.py:274 — not affected: the write timing
metric fires after :340 sets `results_key`, and the key is tracked as well
(write_timing tests pass)
- superset/sql/execution/celery_task.py:707 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (a stop at the publish check is
deleted at :708; if the check raises, the handler has the tracked key)
- superset/sql/execution/celery_task.py:713 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (a guarded UPDATE that raises goes to
the handler with the tracked key, see the publish_update tests; a lost guard is
deleted at :714; a flush failure in its commit is handled by the rollback at
:453, see the publish_flush tests)
- superset/sql/execution/celery_task.py:701 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (when the durable commit raises, the
handler commits FAILED at :127 with the key written back at :116)
- superset/sql/execution/celery_task.py:105 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (the handler's stop check deletes the
tracked key at :106)
- superset/sql/execution/celery_task.py:128 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (the guarded error commit losing to a
stop deletes the tracked key at :129)
- superset/sql/execution/celery_task.py:439 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (an unrecorded failure deletes the
tracked key at :449)
- superset/sql/execution/celery_task.py:322 — not affected: the write failed
and nothing was stored
- superset/sql_lab.py:906 — not affected: the key is tracked at :908 with
only an attribute set in between
- superset/sql_lab.py:930 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (if the final flush or the refresh at
:931 raises, `get_query` rolls back, and the handler deletes the blob on a stop
at :150 or writes the key back on FAILED at :160)
- superset/sql_lab.py:936 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (same as above; after it succeeds,
:938 clears the tracked key because the row references it)
- superset/sql_lab.py:318 — fixed in
88a85e28e8047bc2d3029e1e786285360f8160e7 (an unrecorded failure deletes the
tracked key at :331)
- superset/sql_lab.py:854 — not affected: the write failed and nothing was
stored
Remaining:
- If the worker dies between the write and the publish (hard time limit, OOM
kill), no handler runs, and the blob expires with `cache_timeout`.
- If the handler raises after its own FAILED commit, the blob is deleted
while the row still references it. Only payload building and config reads run
after that commit.
- A failed delete is logged, and the blob expires with `cache_timeout`.
--
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]
