aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4276788902
##########
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: the earlier fix missed one return. If something fails after a
successful write (for example the write timing metric or the guarded publish
UPDATE) and a stop lands at the same time, the task error handler
`_handle_query_error` returned STOPPED and dropped the pending key without
deleting the blob. Fixed in 0f9af0072897bbc8e8e16f35c971d43dae3b1b49: the
handler reads the stored key before its stop check and deletes the blob on both
of its stop returns (the stop check, and the guarded commit finding the row
STOPPED). A failed delete is logged and the query still ends STOPPED.
The legacy path had the same gap, already on master: a failure between the
write and the final metadata commit reaches `handle_query_error`, which returns
STOPPED without committing, so the key is lost. Fixed in
e7ebebc34b8c0281625d79b83d9ac5c072470e6c: `execute_sql_statements` reports the
written key and clears it once the final commit publishes it, and
`handle_query_error` deletes a key that is still unpublished when a stop wins.
A key the query row already references is never deleted.
Regression tests:
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_error_after_results_write_deletes_stored_results_when_stopped`
makes the write timing or the guarded publish UPDATE raise, with the stop
committed from another session either with the failure or after the handler's
stop check, and a working and a failing delete. The 8 stop cases fail at
ba1fcc6b893a15a33f36c33fec060274b41e840c ("Expected 'delete' to be called once.
Called 0 times.") and pass now. The 4 no-stop cases pass on both: the query
ends FAILED, keeps its `results_key`, and nothing is deleted.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_error_after_write_deletes_results_when_stopped`
does the same for the legacy path. The 4 stop cases fail at
0f9af0072897bbc8e8e16f35c971d43dae3b1b49 and pass now; the 4 no-stop cases pass
on both.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_keeps_published_results_of_stopped_query`
is a control: a failure after the final commit of a stopped query keeps the
referenced blob.
Siblings checked:
- superset/sql/execution/celery_task.py:102 — fixed in
0f9af0072897bbc8e8e16f35c971d43dae3b1b49 (stop check in the task error handler).
- superset/sql/execution/celery_task.py:123 — fixed in
0f9af0072897bbc8e8e16f35c971d43dae3b1b49 (guarded error commit finds the row
STOPPED).
- superset/sql_lab.py:149 — fixed in
e7ebebc34b8c0281625d79b83d9ac5c072470e6c (legacy error handler stop return).
- superset/sql/execution/celery_task.py:299 — fixed in
ba1fcc6b893a15a33f36c33fec060274b41e840c (post-write stop check; its test still
passes).
- superset/sql/execution/celery_task.py:659 — fixed in
ba1fcc6b893a15a33f36c33fec060274b41e840c (publish stop check; its test still
passes).
- superset/sql/execution/celery_task.py:665 — fixed in
ba1fcc6b893a15a33f36c33fec060274b41e840c (guarded publish commit; its test
still passes).
- superset/sql/execution/celery_task.py:308 — not affected: the write
failed, so nothing was stored.
- superset/sql_lab.py:840 — not affected: the write failed, so nothing was
stored.
- superset/sql/execution/celery_task.py:120 — not affected: committed
mutation/CTAS work ignores a later stop, so the error is reported FAILED and
the row keeps its key.
- superset/sql/execution/celery_task.py:648 — not affected: durable
publication ignores a later stop and keeps its published results.
Remaining: none for a stopped query. A correction to my earlier reply: a
failure without a stop does not strand the blob. The FAILED row keeps
`results_key` pointing at it (asserted in both new tests), and it expires with
`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]
