EnxDev commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4280569777
##########
superset/sql/execution/celery_task.py:
##########
@@ -360,27 +431,203 @@ def execute_sql_task(
query_id,
rendered_query,
start_time=start_time,
+ on_commit=mark_committed,
+ on_results_key=track_results_key,
),
can_retry=lambda: not query.progress,
)
except Exception as ex:
logger.exception("Query %d: %s", query_id, ex)
stats_logger = app.config["STATS_LOGGER"]
stats_logger.incr("error_sqllab_unhandled")
- query = _get_query(query_id=query_id)
- return _handle_query_error(ex, query)
+ try:
+ query = _get_query_after_error(query_id)
+ # Failures after the target commit (result serialization,
+ # the results backend, closing the connection) cannot undo
+ # it.
+ return _handle_query_error(
+ ex, query, durable=committed, results_key=results_key
+ )
+ except Exception:
+ # _handle_query_error commits last, so the failure was most
+ # likely not recorded. A commit can still apply and raise.
+ _discard_unreferenced_results(query_id, results_key)
+ raise
+
+
+def _get_query_after_error(query_id: int) -> Query:
+ """Get the query, rolling back a session a failed flush left unusable."""
+ try:
+ return _get_query(query_id=query_id)
+ except Exception: # pylint: disable=broad-except
+ # Rolling back only discards metadata the failed task never committed.
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ return _get_query(query_id=query_id)
+
+
+def _commit_query_unless_stopped(query: Query) -> bool:
+ """Atomically guard terminal metadata writes against a concurrent stop."""
+ with db.session.no_autoflush:
+ # The conditional write acquires the row's write lock until commit.
+ # A separate status read cannot protect the subsequent ORM flush.
+ updated = (
+ db.session.query(Query)
+ .filter(Query.id == query.id, Query.status != QueryStatus.STOPPED)
Review Comment:
This guards the worker against a stop, but the stop side has no matching
guard. `QueryDAO.stop_query` (`superset/daos/query.py:73-92`) reads RUNNING,
calls `cancel_query` (a round trip to the engine), and only then flushes
`status=STOPPED` with a plain `WHERE id = ?`. If this UPDATE commits SUCCESS
while that cancel is in flight, the stop's commit lands after it and overwrites
it. The row ends STOPPED with a published `results_key`, and on the `durable`
path that hides a committed INSERT or CTAS, which is the case this PR sets out
to prevent.
Could `stop_query` write STOPPED with the mirror filter (`status NOT IN
(success, failed, timed_out)`), or re-read with `with_for_update` before
setting it? A second-session test that commits SUCCESS between the stop's read
and its commit would pin it down.
##########
superset/sql_lab.py:
##########
@@ -62,6 +63,7 @@
from superset.extensions import celery_app, event_logger
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
+from superset.sql.execution.celery_task import stop_seen_during_execution
Review Comment:
Nit, take it or leave it. The legacy module now imports from the Celery task
module, and it also keeps its own copies of `_extract_errors`,
`_discard_unreferenced_results` and the delete helper, which are identical to
the ones in `celery_task.py`.
The unreferenced-blob rule changed a few times over this PR. Having two
copies means the next fix has to land twice. Maybe move these and
`stop_seen_during_execution` into a small neutral module (or `executor.py`)
that both paths import?
--
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]
