aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4279066634
##########
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 found a case my last reply listed as not
affected (the two publish-stop deletes,
superset/sql/execution/celery_task.py:762 and :768 at
c9abbe220cb582bc28ad51b76be4ce93c633dc0f). When Celery redelivers a task whose
row already committed a `results_key` (for example with `task_acks_late`), and
a stop lands during the replay's results write, `_store_results_in_backend`
deleted the replay's new blob and reloaded the row, which restored the earlier
attempt's key. `_publish_query_result` then read that restored key and deleted
it too, so the STOPPED row referenced deleted results.
Fixed in ea6168068bb2473e25fa39e4c9ccaaab143039fa.
`_store_results_in_backend` returns the key it wrote and left pending (None
when a stop won and it already deleted the blob), and both stop branches of
`_publish_query_result` delete only that key. Publication never deletes the
row's own `results_key`.
Regression test:
`tests/unit_tests/sql/execution/test_celery_task.py::test_replayed_task_stop_keeps_results_the_committed_row_references`
fails at c9abbe220cb582bc28ad51b76be4ce93c633dc0f with `assert [] ==
['prior-attempt-key']` (both blobs deleted) and passes at
ea6168068bb2473e25fa39e4c9ccaaab143039fa. It starts from a committed SUCCESS
row with `results_key='prior-attempt-key'`, runs the full task, commits a stop
from a second session during the replay's write, and asserts the payload and
the persisted row are STOPPED, the row still references `prior-attempt-key`,
that blob is still stored, and only the replay's own blob was deleted.
Siblings checked:
- superset/sql/execution/celery_task.py:814 — fixed in
ea6168068bb2473e25fa39e4c9ccaaab143039fa (stop seen before the publish commit)
- superset/sql/execution/celery_task.py:820 — fixed in
ea6168068bb2473e25fa39e4c9ccaaab143039fa (guarded publish matched no row)
- superset/sql/execution/celery_task.py:317 — not affected: deletes only the
key generated in that call
- superset/sql/execution/celery_task.py:106 — not affected: on a replay
without a write of its own the key is the row's committed one, and
`_discard_unreferenced_results` reads the row in a new transaction and keeps
results it references (checked with a replay whose execution fails after a
stop: STOPPED, no delete, the earlier blob kept)
- superset/sql/execution/celery_task.py:137 — not affected: same
reference-aware cleanup as line 106
- superset/sql/execution/celery_task.py:454 — not affected: deletes only the
key tracked from this attempt's write, and only if no row references it
- superset/sql_lab.py:981 — not affected: `stored_key` is set only from this
attempt's successful write (line 956), never from the loaded row
- superset/sql_lab.py:154 — not affected: same reference-aware cleanup as
superset/sql/execution/celery_task.py:106
- superset/sql_lab.py:370 — not affected: same reference-aware cleanup as
superset/sql/execution/celery_task.py:106
Remaining:
- If the metadata database cannot be read after the failure or the stop, 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.
--
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]
