aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4278245431
##########
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 missed one case and listed it as remaining. In
legacy `get_sql_results`, `handle_query_error` committed FAILED with the
results key and then read `query.status` back for the payload. That read hits
the metadata database again, so when it raised, the outer cleanup treated the
failure as unrecorded and deleted the blob the committed row pointed to. This
happened with `return_results=True` and with `return_results=False`.
Fixed in 24bca1e910010632f50dc19c477fca0ef94625c3:
- `handle_query_error` (superset/sql_lab.py) builds the whole payload before
it commits, using `QueryStatus.FAILED` instead of re-reading the row, so the
commit is the last step that can fail. If the handler raises, the FAILED row
was not committed and nothing references the blob, so the cleanup in
`get_sql_results` deletes it. Once FAILED is committed, nothing after it can
raise into that cleanup.
- `_handle_query_error` (superset/sql/execution/celery_task.py) does the
same: the error payload, stacktrace and troubleshooting link are built before
the plain or guarded commit.
- Timeouts: a stop recorded while the soft time limit fires no longer
replaces TIMED_OUT in the Celery path. The user did not cause the timeout and
it is a real terminal outcome, so like a durable commit it is written without
the stop check. This matches https://github.com/apache/superset/pull/45232. I
found no case where the stop should win: the stop would only relabel a query
the time limit had already killed.
Regression tests (each fails at 88a85e28e8047bc2d3029e1e786285360f8160e7 and
passes at 24bca1e910010632f50dc19c477fca0ef94625c3):
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_deletes_results_only_without_committed_failure[after_failed_commit-False]`
and `[after_failed_commit-True]`: once the FAILED row is committed, every
later metadata statement raises. At 88a85e28 both raised `RuntimeError:
Metadata database unavailable` from the status read at superset/sql_lab.py:173
and deleted the blob. Now the task returns the FAILED payload, and a separate
session sees FAILED with the key while the blob is still stored and `delete` is
never called.
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_deletes_results_only_without_committed_failure[failed_commit-*]`:
the database commit of the FAILED row raises. The row stays uncommitted and
the blob is deleted, in both modes. These cases also passed at 88a85e28 and
guard the other half of the rule.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_error_path_deletes_results_only_without_committed_failure[payload]`:
building the error payload raises. At 88a85e28 it failed with `assert 'failed'
== <QueryStatus.RUNNING>`: FAILED had been committed with the key and the blob
was deleted. Now nothing is committed and the blob is deleted. The `[commit]`
case covers a failed FAILED commit.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_task_error_lookup_does_not_flush_over_a_stop[True-timeout]`
and `[True-timeout_oauth2_lookup]`: a stop is committed in a separate session
while the soft time limit fires. At 88a85e28 both failed with `assert 'stopped'
== <QueryStatus.TIMED_OUT>`. Now the row and the payload are TIMED_OUT.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_timeout_is_recorded_without_reading_a_stop`:
at 88a85e28 it failed with "Expected 'refresh' to not have been called. Called
1 times."
- `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`: 199 passed, 2 skipped.
`superset/sql/` coverage is 100% under `tests/unit_tests/sql/`.
The rule for every delete site: delete a stored blob only when no committed
row references it.
Siblings checked:
- superset/sql_lab.py:335 — fixed in
24bca1e910010632f50dc19c477fca0ef94625c3 (only reached when `get_query` or
`handle_query_error` raises, and the handler's commit at :182 is its last
fallible step)
- superset/sql_lab.py:150 — not affected: the stop path deletes before
returning and commits nothing. The STOPPED row was written by the stop request,
which never sets `results_key`
- superset/sql_lab.py:940 — not affected: the success commit that publishes
the key is followed directly by :942 clearing the tracked key, so later
failures do not delete it
- superset/sql_lab.py:898 — not affected: the write failed, no key is tracked
- superset/sql/execution/celery_task.py:450 — fixed in
24bca1e910010632f50dc19c477fca0ef94625c3 (only reached when
`_get_query_after_error` or `_handle_query_error` raises, and the handler's
commit at :134 or :478 is its last fallible step)
- superset/sql/execution/celery_task.py:105 — not affected: a stop won
before anything committed the key; the stop check reloads the row and drops the
pending key, so the STOPPED commit at :515 carries none
- superset/sql/execution/celery_task.py:136 — not affected: the lost guarded
commit reloads the row and drops the pending key, so the STOPPED commit at :515
carries none
- superset/sql/execution/celery_task.py:314 — not affected: deleted before
any commit references the key
- superset/sql/execution/celery_task.py:710 — not affected: the publish has
not committed the key when the stop check wins
- superset/sql/execution/celery_task.py:716 — not affected: the guarded
publish lost to a stop and committed nothing
- superset/sql/execution/celery_task.py:703 — not affected: the durable
publish commit is the task's last metadata write, so no later failure reaches
:450
- superset/sql/execution/celery_task.py:715 — not affected: the guarded
publish commit is the task's last metadata write, so no later failure reaches
:450
- superset/sql/execution/celery_task.py:330 — not affected: the write failed
and the FAILED row is committed with `results_key=None`
Remaining:
- If a commit raises after the database has applied it (the acknowledgement
is lost), the row is treated as uncommitted and the blob is deleted. The two
cannot be told apart from the worker.
- If `_stopped_query_payload`'s commit raises after a stop-path delete, :450
deletes the same key again. The second delete only logs "Unpublished results
were not deleted".
- The legacy `handle_query_error` still lets a recorded stop replace a
timeout, and it records FAILED rather than TIMED_OUT. That is master's
behaviour (see its TODO) and this PR does not change it.
- Worker death between the write and the publish, and a failed delete, still
leave the blob until `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]
