aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4276440285
##########
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:
Thanks, good catch. Fixed in dd5670cc8185700a68384ee8171ae0940c64a20c: the
sync path now reads the status right before `conn.commit()`, like
`_execute_sql_statements`. A stop recorded while the final statement runs skips
the commit, so the transaction is dropped with the connection and the query
ends STOPPED with no rows written. Only a stop that lands during or after the
commit loses to the committed work and ends SUCCESS.
Regression test:
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_stop_before_target_commit_stays_stopped`
(mutation and CTAS cases) commits a stop from another session while the last
statement runs. It fails at d117709fde (commit runs, query ends SUCCESS) and
passes now. The finish-before-stop case is
`test_get_sql_results_keeps_committed_work_when_stop_races_commit`, which still
passes (stop during the commit ends SUCCESS).
Siblings checked:
- superset/sql/execution/celery_task.py:605 — not affected: the Celery path
already checks for a stop before its target commit and skips it.
- superset/sql_lab.py:608 — not affected: the per-block stop check runs
before each statement starts, so it cannot see a stop recorded during the last
one.
- superset/sql_lab.py:707 — not affected: the post-execution check still
ends STOPPED unless the target commit ran.
- superset/sql_lab.py:815 — not affected: the results-backend write-failure
check still ends STOPPED unless the target commit ran.
Remaining: engines that autocommit each statement (Impala, Hive, Trino) make
the rows permanent when the statement finishes and `conn.commit()` does
nothing, so a stop there ends STOPPED with the rows already written. This is
the same on both paths.
##########
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:
Thanks, I took it. Fixed in ba1fcc6b893a15a33f36c33fec060274b41e840c: when
this check fires after a successful write, the task now calls
`results_backend.delete(key)` before returning. The same happens when the stop
is caught later, at the publish check or when the guarded publish commit finds
the row STOPPED, since the blob is stranded the same way there. If the delete
raises or returns false, it is logged as a warning and the query still ends
STOPPED. A failed write stores nothing, so nothing is deleted.
Regression test:
`tests/unit_tests/sql/execution/test_celery_task.py::test_stop_after_results_write_deletes_stored_results`
commits a stop from another session during the write, after the post-write
check, and after the publish check, with a working and a failing delete. The
six stop-after-successful-write cases fail at
dd5670cc8185700a68384ee8171ae0940c64a20c (delete never called, blob left in the
backend) and pass now. The no-stop and failed-write cases pass on both and
check that nothing is deleted.
Siblings checked:
- superset/sql/execution/celery_task.py:654 — fixed in
ba1fcc6b893a15a33f36c33fec060274b41e840c (stop caught at the publish check
after a successful write).
- superset/sql/execution/celery_task.py:660 — fixed in
ba1fcc6b893a15a33f36c33fec060274b41e840c (guarded publish commit finds the row
STOPPED).
- superset/sql/execution/celery_task.py:296 — not affected: the
write-failure branch stored nothing.
- superset/sql_lab.py:784 — not affected: the legacy path has no stop check
between a successful write and its final commit. A stop during the write keeps
the row STOPPED and commits `results_key`, so the blob is still referenced by
the query, not stranded. That is unchanged from master.
Remaining: an exception after a successful write that is not a stop (for
example the metadata commit in the publish step failing) ends FAILED and still
leaves the blob until `cache_timeout`. This PR does not change that.
--
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]
