You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2026/10/15 23:04:13 UTC
(superset) branch master updated: fix(sql): keep stop and finish outcomes consistent when they race (#45283)
This is an automated email from the ASF dual-hosted git repository.
aminghadersohi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 13bd0cd9326 fix(sql): keep stop and finish outcomes consistent when they race (#45283)
13bd0cd9326 is described below
commit 13bd0cd932632522ddc28d193a1057f8f487691c
Author: Amin Ghadersohi <am...@gmail.com>
AuthorDate: Fri Oct 16 10:04:06 2026 +1100
fix(sql): keep stop and finish outcomes consistent when they race (#45283)
---
UPDATING.md | 20 +
docs/docs/using-superset/exploring-data.mdx | 36 +
docs/static/resources/openapi.json | 3 +
.../src/SqlLab/actions/sqlLab.test.ts | 57 +
superset-frontend/src/SqlLab/actions/sqlLab.ts | 32 +-
superset/constants.py | 1 +
superset/daos/query.py | 26 +-
superset/db_engine_specs/README.md | 2 +
superset/db_engine_specs/base.py | 3 +
superset/db_engine_specs/impala.py | 1 -
superset/db_engine_specs/trino.py | 2 +-
superset/queries/api.py | 15 +-
superset/sql/execution/celery_task.py | 466 ++--
superset/sql/execution/executor.py | 19 +-
superset/sql/execution/utils.py | 260 +++
superset/sql_lab.py | 356 ++--
tests/integration_tests/queries/api_tests.py | 29 +-
tests/integration_tests/sqllab_tests.py | 29 +-
tests/testcontainers/db_engine_specs/test_mysql.py | 113 +
.../db_engine_specs/test_postgres.py | 690 ++++++
tests/unit_tests/dao/queries_test.py | 30 +-
tests/unit_tests/dao/test_query_stop_race.py | 73 +-
tests/unit_tests/daos/test_query_dao.py | 608 ++++++
.../db_engine_specs/test_impala_pending.py | 3 +-
tests/unit_tests/sql/execution/conftest.py | 1 +
tests/unit_tests/sql/execution/test_celery_task.py | 2219 +++++++++++++++++++-
tests/unit_tests/sql/execution/test_executor.py | 3 +-
tests/unit_tests/sql/execution/test_utils.py | 207 ++
tests/unit_tests/sql_lab_test.py | 1714 +++++++++++++++
tests/unit_tests/test_sqllab_error_stacktrace.py | 17 +-
30 files changed, 6575 insertions(+), 460 deletions(-)
diff --git a/UPDATING.md b/UPDATING.md
index acdd6f2b983..5f411684113 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -24,6 +24,26 @@ assists people when migrating to a new version.
## Next
+- `POST /api/v1/query/stop` returns HTTP 422 instead of 200 when the query
+ has already finished or its worker has committed the target transaction.
+ The response contains `message: "Could not cancel query"` and the recorded
+ `query`; an async handle's `cancel()` returns `False`. After a target commit,
+ that query can still be running while SQL Lab prepares its results.
+ PostgreSQL/MySQL metadata databases serialize Stop against the worker's final
+ stop check and target commit. If Stop wins, transactional work is discarded;
+ if the worker wins, Stop is rejected even before results are published.
+ An accepted Stop returns HTTP 200 (`cancel()` returns `True`) even if engine
+ cancellation or subsequent cancellation-metadata persistence fails. Its stopped
+ state and original end time survive repeated Stop requests and worker timeouts.
+ Repeated Stop requests do not retry engine cancellation.
+ Engine cancellation is best-effort: its statement may continue, but SQL Lab
+ discards the results. Stop cannot undo autocommit, implicit-commit DDL, or SQL
+ commits performed during execution. With SQLite metadata (which does not
+ support row locks), a Stop accepted between the worker's final check and
+ target commit can also leave transactional changes committed. Target and
+ metadata commits are not atomic: a metadata commit failure after the target
+ commit can likewise leave database changes without the committed marker.
+
- **Semantic NULL filters and provider SDK upgrade.** Semantic filters normalize
UI NULL/empty-string selections. Mixed positive NULL membership uses `OrFilter`,
which every semantic provider must render; previously NULL rows could be
diff --git a/docs/docs/using-superset/exploring-data.mdx b/docs/docs/using-superset/exploring-data.mdx
index 444f0ae5bd1..8084fda4b0f 100644
--- a/docs/docs/using-superset/exploring-data.mdx
+++ b/docs/docs/using-superset/exploring-data.mdx
@@ -24,6 +24,40 @@ The following information about each flight is given:
- Information about the origin and destination.
- The distance between the origin and destination, in kilometers (km).
+### Stopping a SQL Lab query
+
+SQL Lab records Stop before attempting engine cancellation. Once recorded, the
+Stop API returns HTTP 200, the UI shows the query as stopped, and an async query
+handle's `cancel()` returns `True`. Cancellation-related failures, later worker
+timeouts, and failures to save cancellation metadata do not replace that outcome.
+Repeated Stop requests preserve the original stop time and do not retry engine
+cancellation.
+
+Engine cancellation is best-effort: if the engine is unreachable, provides no
+cancellation handle, or rejects cancellation, its statement may run to completion,
+but SQL Lab discards its results. The UI still shows “Query was stopped.”, not a
+stop-failure toast. Stopping cannot undo changes the engine commits during
+execution, including autocommit mutations, implicit-commit DDL, or explicit SQL
+commits.
+
+With PostgreSQL or MySQL metadata databases, Stop and the worker's final target
+transaction commit are serialized using a query-row lock. If Stop wins, the
+worker does not commit its transactional changes. If the worker wins, Stop is
+rejected with HTTP 422 (`cancel()` returns `False`), even while SQL Lab is still
+preparing the results. The response contains the recorded query, which can still
+be running at that point. A query that already finished likewise keeps its
+success, failure, or timeout outcome. SQL Lab refreshes a finished query's outcome
+without a stop-failure toast. If the target commit wins while the query is still
+running, SQL Lab shows “Failed to stop query.” instead: the changes have committed,
+but result preparation has not finished.
+
+SQLite metadata databases do not support this row lock: a Stop accepted between
+the worker's final stop check and its target commit can still leave transactional
+changes committed. Target and metadata commits cannot be atomic across databases;
+if recording the target-committed marker fails after the target commit, changes
+can also persist without that marker. A stopped state is therefore not a general
+promise that database changes were rolled back.
+
### Semantic Table time grains
For a Table chart backed by a semantic view, Time Grain applies when grouping by
@@ -411,6 +445,8 @@ tutorial dashboard to see the four charts side by side and compare the different
### SQL Lab Tips
+**Stopping queries**: An accepted Stop shows “Query was stopped.” even if engine cancellation fails. A finished query refreshes its completed outcome without a stop-failure toast. A Stop rejected after the target commit, while result preparation is still running, shows the “Failed to stop query.” toast. See [Stopping a SQL Lab query](#stopping-a-sql-lab-query) for transaction guarantees and metadata-database limitations.
+
**Schema and table browser**: The left-side table browser uses a collapsible treeview — click a schema to expand its tables, and click a table to see its columns and sample data inline. This makes navigating large schemas much faster than the previous flat list.
**Find in editor**: Press **Ctrl+F** (or **Cmd+F** on Mac) to open the Monaco find/replace widget inside the SQL editor without leaving the editor.
diff --git a/docs/static/resources/openapi.json b/docs/static/resources/openapi.json
index b07196c19a9..ad86fd1a91e 100644
--- a/docs/static/resources/openapi.json
+++ b/docs/static/resources/openapi.json
@@ -27375,6 +27375,9 @@
"404": {
"$ref": "#/components/responses/404"
},
+ "422": {
+ "description": "Query already finished or target transaction committed; includes the recorded query\\n"
+ },
"500": {
"$ref": "#/components/responses/500"
}
diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.ts b/superset-frontend/src/SqlLab/actions/sqlLab.test.ts
index 28b20042dcc..c524b641a46 100644
--- a/superset-frontend/src/SqlLab/actions/sqlLab.test.ts
+++ b/superset-frontend/src/SqlLab/actions/sqlLab.test.ts
@@ -2212,3 +2212,60 @@ describe('async actions', () => {
});
});
});
+
+test.each(['success', 'failed', 'timed_out'])(
+ 'postStopQuery refreshes an already %s query without a toast',
+ async state => {
+ const finishedQuery = { ...queryFixture, id: queryId, state };
+ const post = jest.spyOn(SupersetClient, 'post').mockRejectedValueOnce({
+ status: 422,
+ json: async () => ({
+ message: 'Could not cancel query',
+ query: finishedQuery,
+ }),
+ });
+ const store = mockStore(initialState);
+ await store.dispatch(
+ actions.postStopQuery({ ...queryFixture, id: queryId }),
+ );
+ expect(store.getActions()).toEqual([
+ actions.refreshQueries({ [finishedQuery.id]: finishedQuery }),
+ ]);
+ post.mockRestore();
+ },
+);
+
+test('postStopQuery retains the error toast for an unsupported stop', async () => {
+ const post = jest.spyOn(SupersetClient, 'post').mockRejectedValueOnce({
+ status: 422,
+ json: async () => ({ message: 'Could not cancel query' }),
+ });
+ const store = mockStore(initialState);
+ await store.dispatch(actions.postStopQuery({ ...queryFixture, id: queryId }));
+ expect(store.getActions()).toEqual([
+ expect.objectContaining({
+ type: ADD_TOAST,
+ payload: expect.objectContaining({ toastType: ToastType.Danger }),
+ }),
+ ]);
+ post.mockRestore();
+});
+
+test('postStopQuery keeps the error toast when a target commit wins before publication', async () => {
+ const post = jest.spyOn(SupersetClient, 'post').mockRejectedValueOnce({
+ status: 422,
+ json: async () => ({
+ message: 'Could not cancel query',
+ query: { ...queryFixture, id: queryId, state: 'running' },
+ }),
+ });
+ const store = mockStore(initialState);
+ await store.dispatch(actions.postStopQuery({ ...queryFixture, id: queryId }));
+ expect(store.getActions()).toEqual([
+ expect.objectContaining({
+ type: ADD_TOAST,
+ payload: expect.objectContaining({ toastType: ToastType.Danger }),
+ }),
+ ]);
+ post.mockRestore();
+});
diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.ts b/superset-frontend/src/SqlLab/actions/sqlLab.ts
index 5fb01e15150..883f6a8d457 100644
--- a/superset-frontend/src/SqlLab/actions/sqlLab.ts
+++ b/superset-frontend/src/SqlLab/actions/sqlLab.ts
@@ -603,6 +603,12 @@ export function reRunQuery(query: Query): SqlLabThunkAction {
};
}
+export function refreshQueries(
+ alteredQueries: Record<string, Query>,
+): SqlLabAction {
+ return { type: REFRESH_QUERIES, alteredQueries };
+}
+
export function postStopQuery(
query: Query,
): SqlLabThunkAction<Promise<unknown>> {
@@ -614,7 +620,25 @@ export function postStopQuery(
})
.then(() => dispatch(stopQuery(query)))
.then(() => dispatch(addSuccessToast(t('Query was stopped.'))))
- .catch(() => dispatch(addDangerToast(t('Failed to stop query.'))));
+ .catch(async (response: Response) => {
+ if (response.status === 422) {
+ const payload: { query?: Query } = await response
+ .json()
+ .catch(() => ({}));
+ const finishedQuery = payload.query;
+ if (
+ finishedQuery &&
+ finishedQuery.id === query.id &&
+ ['success', 'failed', 'timed_out'].includes(
+ finishedQuery.state ?? '',
+ )
+ ) {
+ dispatch(refreshQueries({ [query.id]: finishedQuery }));
+ return;
+ }
+ }
+ dispatch(addDangerToast(t('Failed to stop query.')));
+ });
};
}
@@ -1519,12 +1543,6 @@ export function removeTables(
};
}
-export function refreshQueries(
- alteredQueries: Record<string, Query>,
-): SqlLabAction {
- return { type: REFRESH_QUERIES, alteredQueries };
-}
-
export function setUserOffline(offline: boolean): SqlLabAction {
return { type: SET_USER_OFFLINE, offline };
}
diff --git a/superset/constants.py b/superset/constants.py
index db31281aedb..818683352df 100644
--- a/superset/constants.py
+++ b/superset/constants.py
@@ -62,6 +62,7 @@ QUERY_EARLY_CANCEL_KEY = "early_cancel_query"
# fabricate a stop) apart from "this engine just has no cancel support"
# (must fail honestly) when no cancel ID is on record.
QUERY_DISPATCHED_KEY = "query_dispatched"
+QUERY_TARGET_COMMITTED_KEY = "target_committed"
LRU_CACHE_MAX_SIZE = 256
diff --git a/superset/daos/query.py b/superset/daos/query.py
index 5b4a4b7f186..1ebce6df53a 100644
--- a/superset/daos/query.py
+++ b/superset/daos/query.py
@@ -21,14 +21,14 @@ from typing import Any, Union
from superset import sql_lab
from superset.common.db_query_status import QueryStatus
from superset.daos.base import BaseDAO
-from superset.exceptions import QueryNotFoundException, SupersetCancelQueryException
+from superset.exceptions import QueryNotFoundException
from superset.extensions import db
from superset.models.sql_lab import Query, SavedQuery
from superset.queries.filters import QueryFilter
from superset.queries.saved_queries.filters import SavedQueryFilter
+from superset.sql.execution.utils import finish_query_stop, set_query_stopped
from superset.utils.core import get_user_id
from superset.utils.dates import now_as_float
-from superset.utils.decorators import transaction
logger = logging.getLogger(__name__)
@@ -60,8 +60,8 @@ class QueryDAO(BaseDAO[Query]):
)
@staticmethod
- @transaction()
- def stop_query(client_id: str) -> None:
+ def stop_query(client_id: str) -> Query:
+ """Stop a query and return its recorded outcome, including a lost race."""
query = (
db.session.query(Query)
.filter(Query.client_id == client_id, Query.user_id == get_user_id())
@@ -78,18 +78,14 @@ class QueryDAO(BaseDAO[Query]):
logger.warning(
"Query with client_id could not be stopped: query already complete",
)
- return
+ return query
- if not sql_lab.cancel_query(query):
- raise SupersetCancelQueryException("Could not cancel query")
-
- # cancel_query() may have staged an early-cancel flag on query.extra
- # without committing it (see its docstring/comments); the
- # @transaction decorator commits it together with status=STOPPED
- # below in one transaction, closing the window where another
- # request could observe the flag set but the status still RUNNING.
- query.status = QueryStatus.STOPPED
- query.end_time = now_as_float()
+ with db.session.no_autoflush:
+ if not set_query_stopped(query, now_as_float()):
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return query
+ finish_query_stop(query, lambda: sql_lab.cancel_query(query))
+ return query
class SavedQueryDAO(BaseDAO[SavedQuery]):
diff --git a/superset/db_engine_specs/README.md b/superset/db_engine_specs/README.md
index 49af5da5fec..9b20995403e 100644
--- a/superset/db_engine_specs/README.md
+++ b/superset/db_engine_specs/README.md
@@ -1322,6 +1322,8 @@ Some databases have an implicit query cancelation. When a cursor stops being pol
For other databases, DB engine specs can implement query cancelation via the `prepare_cancel_query` and `cancel_query` methods. Implementation of query cancelation is usually heavily dependent on the database, but the DB engine specs that support it can serve as an example.
+These methods must not flush or commit the metadata session. `prepare_cancel_query` can stage cancellation intent on the query; its caller persists that intent together with a conditional stop, preserving any terminal result committed while cancellation was in flight.
+
### Get metrics on dataset creation
When a physical dataset is first created, the `get_metrics` class method is called on the table. The base implementation returns the `COUNT(*)` metric, but DB engine specs can override `get_metrics` to return other metrics. This method is useful for semantic layers that contain their own metrics definitions; when Superset connect to them it can automatically create those metrics when a dataset is added.
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index b7db9e661d6..57d531aa90d 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -3016,6 +3016,9 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
cancelation request has been received. For those cases, the db engine spec
can record the cancelation intent so that the query can either be stopped
prior to execution, or canceled once the query id is acquired.
+
+ Only stage changes on the query; do not flush or commit here. The caller
+ persists them with a guarded stop so a concurrent terminal result wins.
"""
return None
diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py
index 06e1f8e1c82..c40db88cd72 100644
--- a/superset/db_engine_specs/impala.py
+++ b/superset/db_engine_specs/impala.py
@@ -242,7 +242,6 @@ class ImpalaEngineSpec(BaseEngineSpec):
"""
if QUERY_CANCEL_KEY not in query.extra:
query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
- db.session.commit() # pylint: disable=consider-using-transaction
@classmethod
def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]:
diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
index 010f393acf2..a6cd167ec01 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -502,9 +502,9 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
@classmethod
def prepare_cancel_query(cls, query: Query) -> None:
+ """Stage early-cancel intent for the caller's guarded stop transaction."""
if QUERY_CANCEL_KEY not in query.extra:
query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
- db.session.commit() # pylint: disable=consider-using-transaction
@classmethod
def cancel_query(cls, cursor: Cursor, query: Query, cancel_query_id: str) -> bool:
diff --git a/superset/queries/api.py b/superset/queries/api.py
index 93ce2df7883..828aded4ac9 100644
--- a/superset/queries/api.py
+++ b/superset/queries/api.py
@@ -28,10 +28,11 @@ from flask_appbuilder.api import (
from flask_appbuilder.models.sqla.interface import SQLAInterface
from superset import db, event_logger
+from superset.common.db_query_status import QueryStatus
from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
from superset.daos.query import QueryDAO
from superset.databases.filters import DatabaseFilter
-from superset.exceptions import SupersetException
+from superset.exceptions import SupersetCancelQueryException, SupersetException
from superset.models.sql_lab import Query
from superset.queries.filters import QueryFilter
from superset.queries.schemas import (
@@ -275,12 +276,22 @@ class QueryRestApi(BaseSupersetModelRestApi):
$ref: '#/components/responses/401'
404:
$ref: '#/components/responses/404'
+ 422:
+ description: >
+ Query already finished or target transaction committed;
+ includes the recorded query
500:
$ref: '#/components/responses/500'
"""
try:
body = self.stop_query_schema.load(request.json)
- QueryDAO.stop_query(body["client_id"])
+ query = QueryDAO.stop_query(body["client_id"])
+ if query.status != QueryStatus.STOPPED:
+ return self.response(
+ SupersetCancelQueryException.status,
+ message="Could not cancel query",
+ query=query.to_dict(),
+ )
return self.response(200, result="OK")
except SupersetException as ex:
return self.response(ex.status, message=ex.message)
diff --git a/superset/sql/execution/celery_task.py b/superset/sql/execution/celery_task.py
index a2cf7808170..9e5cb0c7037 100644
--- a/superset/sql/execution/celery_task.py
+++ b/superset/sql/execution/celery_task.py
@@ -27,6 +27,7 @@ import dataclasses
import logging
import traceback
import uuid
+from collections.abc import Callable
from typing import Any, TYPE_CHECKING
import msgpack
@@ -40,12 +41,13 @@ from superset import (
security_manager,
)
from superset.common.db_query_status import QueryStatus
-from superset.constants import QUERY_CANCEL_KEY
-from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
-from superset.exceptions import (
- SupersetErrorException,
- SupersetErrorsException,
+from superset.constants import (
+ QUERY_CANCEL_KEY,
+ QUERY_DISPATCHED_KEY,
+ QUERY_EARLY_CANCEL_KEY,
)
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
+from superset.exceptions import SupersetErrorException
from superset.extensions import celery_app
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
@@ -53,6 +55,15 @@ from superset.sql.execution.executor import (
build_statement_blocks,
execute_sql_with_cursor,
)
+from superset.sql.execution.utils import (
+ _discard_stored_results,
+ _discard_unreferenced_results,
+ _extract_errors,
+ commit_query_unless_stopped as _commit_query_unless_stopped,
+ commit_target_transaction,
+ set_query_stopped,
+ stop_seen_during_execution,
+)
from superset.sql.parse import SQLScript
from superset.sqllab.utils import write_ipc_buffer
from superset.utils import json
@@ -78,43 +89,54 @@ def _handle_query_error(
query: Query,
payload: dict[str, Any] | None = None,
prefix_message: str = "",
+ results_key: str | None = None,
) -> dict[str, Any]:
- """Handle error while processing the SQL query."""
- payload = payload or {}
- msg = f"{prefix_message} {str(ex)}".strip()
- query.error_message = msg
- query.tmp_table_name = None
- # Preserve TIMED_OUT status if already set (from SoftTimeLimitExceeded handler)
- if query.status != QueryStatus.TIMED_OUT:
- query.status = QueryStatus.FAILED
-
- if not query.end_time:
- query.end_time = now_as_float()
-
- # Extract DB-specific errors
- if isinstance(ex, SupersetErrorException):
- errors = [ex.error]
- elif isinstance(ex, SupersetErrorsException):
- errors = ex.errors
- else:
- errors = query.database.db_engine_spec.extract_errors(
- str(ex), database_name=query.database.unique_name
- )
-
- errors_payload = [dataclasses.asdict(error) for error in errors]
- if errors:
- query.set_extra_json_key("errors", errors_payload)
+ """
+ Handle error while processing the SQL query.
- db.session.commit() # pylint: disable=consider-using-transaction
- payload.update(
- {"status": query.status.value, "error": msg, "errors": errors_payload}
- )
- if app.config.get("SHOW_STACKTRACE"):
- if stacktrace := traceback.format_exc():
- payload["stacktrace"] = stacktrace
- if troubleshooting_link := app.config.get("TROUBLESHOOTING_LINK"):
- payload["link"] = troubleshooting_link
- return payload
+ :param results_key: key of results stored in the results backend before the
+ failure; the failed row keeps it, or it is deleted if a stop wins
+ """
+ # Preserve an uncommitted timeout only if no accepted Stop won the row.
+ timed_out = query.status == QueryStatus.TIMED_OUT
+ status = QueryStatus.TIMED_OUT if timed_out else QueryStatus.FAILED
+ # Results stored before the failure are unpublished; a stop reloads the row
+ # and drops their pending key, so read it first. A rollback after a failed
+ # flush has already dropped it, so prefer the key tracked at the write.
+ stored_key = results_key or query.results_key
+ if _make_check_stopped_fn(query, preserve_pending=True)():
+ _discard_unreferenced_results(query.id, stored_key, results_backend)
+ return _stopped_query_payload(query)
+
+ with db.session.no_autoflush:
+ payload = payload or {}
+ msg = f"{prefix_message} {str(ex)}".strip()
+ query.error_message = msg
+ query.tmp_table_name = None
+ query.status = status
+ if stored_key:
+ query.results_key = stored_key
+
+ if not query.end_time:
+ query.end_time = now_as_float()
+
+ errors = _extract_errors(ex, query)
+ errors_payload = [dataclasses.asdict(error) for error in errors]
+ if errors:
+ query.set_extra_json_key("errors", errors_payload)
+
+ payload.update({"status": status.value, "error": msg, "errors": errors_payload})
+ if app.config.get("SHOW_STACKTRACE"):
+ if stacktrace := traceback.format_exc():
+ payload["stacktrace"] = stacktrace
+ if troubleshooting_link := app.config.get("TROUBLESHOOTING_LINK"):
+ payload["link"] = troubleshooting_link
+ # The commit is the last step that can fail: when this function raises,
+ # the failed row was not committed and does not reference the results.
+ if not _commit_query_unless_stopped(query):
+ _discard_unreferenced_results(query.id, stored_key, results_backend)
+ return _stopped_query_payload(query)
+ return payload
def _serialize_payload(payload: dict[Any, Any]) -> bytes:
@@ -198,102 +220,120 @@ 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,
-) -> 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,
- )
- )
-
- cache_timeout = database.cache_timeout
- if cache_timeout is None:
- cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"]
+ on_results_key: Callable[[str], None] | None = None,
+) -> str | None:
+ """
+ Store query results in the results backend.
- compressed = zlib_compress(serialized_payload)
- logger.debug("*** serialized payload size: %i", len(serialized_payload))
- logger.debug("*** compressed payload size: %i", len(compressed))
+ :param on_results_key: called with the key as soon as the results are
+ stored, before any metadata read that can fail
+ :returns: the key set on the row and not yet published, or None if a stop
+ won and the stored results were deleted
+ """
+ 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,
+ )
+ )
- 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 write_success and on_results_key:
+ on_results_key(key)
+ if _make_check_stopped_fn(query, preserve_pending=True)():
+ if write_success:
+ _discard_stored_results(query.id, key, results_backend)
+ return None
+ if not write_success:
+ logger.error(
+ "Query %s: Failed to store results in backend, key: %s",
+ str(query.id),
+ key,
)
- )
- else:
- query.results_key = key
- logger.info(
- "Query %s: Successfully stored 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."
+ )
+ if not _commit_query_unless_stopped(query):
+ return None
+ raise SupersetErrorException(
+ SupersetError(
+ message=__("Failed to store query results. Please try again."),
+ error_type=SupersetErrorType.RESULTS_BACKEND_ERROR,
+ level=ErrorLevel.ERROR,
+ )
+ )
+ else:
+ query.results_key = key
+ logger.info(
+ "Query %s: Successfully stored results in backend, key: %s",
+ str(query.id),
+ key,
+ )
+ return key
def _serialize_result_set(
@@ -349,7 +389,18 @@ def execute_sql_task(
:returns: Query result payload or None
"""
with app.test_request_context():
- with override_user(security_manager.find_user(username)):
+ # OAuth retries and error recovery can load ORM state after a failure.
+ # Only guarded commits may flush pending terminal metadata over a stop.
+ with (
+ override_user(security_manager.find_user(username)),
+ db.session.no_autoflush,
+ ):
+ results_key: str | None = None
+
+ def track_results_key(key: str) -> None:
+ nonlocal results_key
+ results_key = key
+
try:
from superset.utils.oauth2 import execute_with_oauth2_retry
@@ -360,6 +411,7 @@ def execute_sql_task(
query_id,
rendered_query,
start_time=start_time,
+ on_results_key=track_results_key,
),
can_retry=lambda: not query.progress,
)
@@ -367,20 +419,91 @@ def execute_sql_task(
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)
+ return _handle_query_error(ex, query, 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, results_backend
+ )
+ raise
-def _make_check_stopped_fn(query: Query) -> Any:
- """Create a function to check if query was stopped."""
+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 _stopped_query_payload(query: Query) -> dict[str, Any]:
+ """Return the recorded outcome without overwriting a completed query."""
+ set_query_stopped(query, query.end_time or now_as_float())
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return {"query_id": query.id, "status": query.status}
+
+
+def _make_check_stopped_fn(
+ query: Query, preserve_pending: bool = False
+) -> Callable[[], bool]:
+ """Check stop intent, optionally preserving uncommitted success metadata."""
def check_stopped() -> bool:
- db.session.refresh(query)
- return query.status == QueryStatus.STOPPED
+ # Finalization has dirty result metadata. Read status without flushing
+ # stale state or discarding that metadata on an ordinary success path.
+ with db.session.no_autoflush:
+ db.session.refresh(
+ query, attribute_names=["status"] if preserve_pending else None
+ )
+ stopped = query.status == QueryStatus.STOPPED or bool(
+ query.extra.get(QUERY_EARLY_CANCEL_KEY)
+ )
+ if stopped and preserve_pending:
+ # Discard pending success fields and reload the recorded stop
+ # time before any subsequent commit can overwrite them.
+ db.session.refresh(query)
+ return stopped
return check_stopped
+class _StopSeenTracker:
+ """
+ Record a stop seen while the statements run.
+
+ The executor commits progress after each statement, which expires the
+ session's copy of the row. A stop the engine loaded while a statement ran,
+ or a check between statements saw, is recorded before that commit.
+ """
+
+ def __init__(
+ self,
+ query: Query,
+ check_stopped: Callable[[], bool],
+ execute: Callable[[Any, str], None],
+ ) -> None:
+ self._query = query
+ self._check_stopped = check_stopped
+ self._execute = execute
+ self.seen = False
+
+ def check_stopped(self) -> bool:
+ """Check for a stop between statements."""
+ if self._check_stopped():
+ self.seen = True
+ return self.seen
+
+ def execute(self, cursor: Any, sql: str) -> None:
+ """Execute a statement, then record a stop the engine loaded."""
+ self._execute(cursor, sql)
+ self.seen = self.seen or stop_seen_during_execution(self._query)
+
+
def _make_execute_fn(query: Query, db_engine_spec: Any) -> Any:
"""Create an execute function with stats timing."""
@@ -414,8 +537,14 @@ def _execute_sql_statements(
query_id: int,
rendered_query: str,
start_time: float | None,
+ on_results_key: Callable[[str], None] | None = None,
) -> dict[str, Any] | None:
- """Execute SQL statements and store results."""
+ """
+ Execute SQL statements and store results.
+
+ :param on_results_key: called with the key of results stored in the results
+ backend, before any metadata read or write that can fail
+ """
if start_time:
stats_logger = app.config["STATS_LOGGER"]
stats_logger.timing("sqllab.query.time_pending", now_as_float() - start_time)
@@ -425,12 +554,18 @@ def _execute_sql_statements(
database = query.database
db_engine_spec = database.db_engine_spec
db_engine_spec.patch()
+ check_stopped = _make_check_stopped_fn(query)
+ stop_tracker = _StopSeenTracker(
+ query, check_stopped, _make_execute_fn(query, db_engine_spec)
+ )
logger.info("Query %s: Set query to 'running'", str(query_id))
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
execution_start_time = now_as_float()
- db.session.commit() # pylint: disable=consider-using-transaction
+ if not _commit_query_unless_stopped(query):
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return {"query_id": query_id, "status": QueryStatus.STOPPED}
# Parse original SQL (from user) to preserve before transformations
original_script = SQLScript(query.sql, engine=db_engine_spec.engine)
@@ -439,6 +574,7 @@ def _execute_sql_statements(
parsed_script, blocks = _prepare_statement_blocks(
rendered_query, db_engine_spec, database
)
+ writes = parsed_script.has_mutation() or query.select_as_cta
with database.get_raw_connection(
catalog=query.catalog,
@@ -447,9 +583,11 @@ def _execute_sql_statements(
cursor = conn.cursor()
cancel_query_id = db_engine_spec.get_cancel_query_id(cursor, query)
+ # Dispatch does not imply the engine provides a cancellation handle.
+ query.set_extra_json_key(QUERY_DISPATCHED_KEY, True)
if cancel_query_id is not None:
query.set_extra_json_key(QUERY_CANCEL_KEY, cancel_query_id)
- db.session.commit() # pylint: disable=consider-using-transaction
+ db.session.commit() # pylint: disable=consider-using-transaction
try:
execution_results = execute_sql_with_cursor(
@@ -458,8 +596,8 @@ def _execute_sql_statements(
statements=blocks,
query=query,
log_query_fn=_make_log_query_fn(database),
- check_stopped_fn=_make_check_stopped_fn(query),
- execute_fn=_make_execute_fn(query, db_engine_spec),
+ check_stopped_fn=stop_tracker.check_stopped,
+ execute_fn=stop_tracker.execute,
# `blocks` is a single un-split block when the engine runs multiple
# statements as one; otherwise each block is an individual statement.
is_split=not db_engine_spec.run_multiple_statements_as_one,
@@ -481,25 +619,59 @@ def _execute_sql_statements(
)
) from ex
- # Check if stopped
- if not execution_results:
- payload.update({"status": QueryStatus.STOPPED.value})
- return payload
+ # Re-read stop intent even on autocommit connections: completed engine
+ # effects cannot be undone, but an accepted Stop still discards results.
+ # Engine polling may have observed a stop before a progress commit
+ # expired the worker's copy, so keep that observation as well.
+ stopped = stop_tracker.seen or not execution_results or check_stopped()
# Commit for mutations
- if parsed_script.has_mutation() or query.select_as_cta:
- conn.commit() # pylint: disable=consider-using-transaction
+ if not stopped and writes:
+ stopped = not commit_target_transaction(query, conn.commit)
+
+ if stopped or (writes and check_stopped()):
+ return _stopped_query_payload(query)
+ # Serialization and backend writes can outlive this check. Publication
+ # uses a guarded write even if the target mutation already committed.
total_execution_time_ms = (now_as_float() - execution_start_time) * 1000
_finalize_successful_query(
query, original_script, execution_results, payload, total_execution_time_ms
)
+ results_key = None
if results_backend:
- _store_results_in_backend(query, payload, database)
+ results_key = _store_results_in_backend(
+ query,
+ payload,
+ database,
+ on_results_key=on_results_key,
+ )
+
+ return _publish_query_result(query, payload, results_key=results_key)
+
+
+def _publish_query_result(
+ query: Query,
+ payload: dict[str, Any],
+ results_key: str | None = None,
+) -> dict[str, Any]:
+ """
+ Publish terminal metadata only if no stop has won the finalization race.
+
+ :param results_key: key of results this task stored and has not published;
+ they are deleted if a stop wins. The row's own key is never deleted: a
+ replayed task can load a key an earlier attempt committed.
+ """
+ # Serialization and backend writes can outlive the last execution check.
+ if _make_check_stopped_fn(query, preserve_pending=True)():
+ _discard_stored_results(query.id, results_key, results_backend)
+ return _stopped_query_payload(query)
if query.status != QueryStatus.FAILED:
query.status = QueryStatus.SUCCESS
- db.session.commit() # pylint: disable=consider-using-transaction
+ if not _commit_query_unless_stopped(query):
+ _discard_stored_results(query.id, results_key, results_backend)
+ return _stopped_query_payload(query)
return payload
diff --git a/superset/sql/execution/executor.py b/superset/sql/execution/executor.py
index d6681503864..ef38537d328 100644
--- a/superset/sql/execution/executor.py
+++ b/superset/sql/execution/executor.py
@@ -1428,14 +1428,27 @@ class SQLExecutor:
@staticmethod
def _cancel_async_query(query_id: int, database: Database) -> bool:
- """Cancel an async query."""
+ """Accept a stop with best-effort engine cancellation, or report a lost race."""
+ from superset.common.db_query_status import QueryStatus as DBQueryStatus
from superset.models.sql_lab import Query as QueryModel
+ from superset.sql.execution.utils import finish_query_stop, set_query_stopped
+ from superset.utils.dates import now_as_float
query = db.session.query(QueryModel).filter_by(id=query_id).one_or_none()
- if not query:
+ if not query or query.status in (
+ DBQueryStatus.SUCCESS,
+ DBQueryStatus.FAILED,
+ DBQueryStatus.TIMED_OUT,
+ ):
return False
- return SQLExecutor._cancel_query(database, query)
+ with db.session.no_autoflush:
+ if not set_query_stopped(query, now_as_float()):
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return False
+ return finish_query_stop(
+ query, lambda: SQLExecutor._cancel_query(database, query)
+ )
@staticmethod
def _cancel_query(database: Database, query: Any) -> bool:
diff --git a/superset/sql/execution/utils.py b/superset/sql/execution/utils.py
new file mode 100644
index 00000000000..5f2cbd19e09
--- /dev/null
+++ b/superset/sql/execution/utils.py
@@ -0,0 +1,260 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Shared SQL execution error, result cleanup, and stop finalization helpers."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import sqlalchemy as sa
+
+from superset import db
+from superset.common.db_query_status import QueryStatus
+from superset.constants import QUERY_EARLY_CANCEL_KEY, QUERY_TARGET_COMMITTED_KEY
+from superset.errors import SupersetError
+from superset.exceptions import SupersetErrorException, SupersetErrorsException
+from superset.models.sql_lab import Query
+from superset.utils import json
+from superset.utils.dates import now_as_float
+
+if TYPE_CHECKING:
+ from flask_caching.backends.base import BaseCache
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_errors(ex: Exception, query: Query) -> list[SupersetError]:
+ """Extract DB-specific errors."""
+ if isinstance(ex, SupersetErrorException):
+ return [ex.error]
+ if isinstance(ex, SupersetErrorsException):
+ return ex.errors
+ return query.database.db_engine_spec.extract_errors(
+ str(ex), database_name=query.database.unique_name
+ )
+
+
+def _discard_stored_results(
+ query_id: int, key: str | None, results_backend: BaseCache | None
+) -> None:
+ """
+ Delete stored results whose key no committed query row will reference.
+
+ That happens when a stop wins, or when the task cannot record its failure.
+ The blob would only occupy the results backend until it expires. A failed
+ delete is logged and does not change the query's outcome.
+ """
+ if not key or not results_backend:
+ return
+ try:
+ deleted = results_backend.delete(key)
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Query %s: Failed to delete unpublished results, key: %s",
+ str(query_id),
+ key,
+ exc_info=True,
+ )
+ return
+ if not deleted:
+ logger.warning(
+ "Query %s: Unpublished results were not deleted, key: %s",
+ str(query_id),
+ key,
+ )
+
+
+def _discard_unreferenced_results(
+ query_id: int, key: str | None, results_backend: BaseCache | None
+) -> None:
+ """
+ Delete stored results after a failure the task could not record, or after
+ a stop wins the error path, unless the query row references them.
+
+ A metadata commit can apply and still raise, for example when its
+ acknowledgement is lost, and a stop can then land on the row that commit
+ published. Uncommitted metadata is rolled back and the row is read again
+ in a new transaction. If that read fails, the results are kept until they
+ expire: an orphaned blob is safer than a committed row that references
+ deleted results.
+ """
+ if not key or not results_backend:
+ return
+ try:
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ referenced_key = (
+ db.session.query(Query.results_key).filter(Query.id == query_id).scalar()
+ )
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Query %s: Kept results that may be referenced, key: %s",
+ str(query_id),
+ key,
+ exc_info=True,
+ )
+ return
+ if referenced_key != key:
+ _discard_stored_results(query_id, key, results_backend)
+
+
+def stop_seen_during_execution(query: Query) -> bool:
+ """
+ Whether the session's copy of the query held a stop when execution returned.
+
+ Engines that cancel a running statement on a stop (Hive, Impala, Presto) do
+ so after loading the stop into this session, as do the checks between
+ statements. Attributes expired by a commit count as unseen. This local
+ observation cannot rule out a stop committed by another session; callers
+ must also check the recorded state before publishing results.
+ """
+ attrs = sa.inspect(query).attrs
+ if attrs.status.loaded_value == QueryStatus.STOPPED:
+ return True
+ extra_json = attrs.extra_json.loaded_value
+ return bool(
+ isinstance(extra_json, str)
+ and extra_json
+ and json.loads(extra_json).get(QUERY_EARLY_CANCEL_KEY)
+ )
+
+
+def set_query_stopped(query: Query, end_time: float) -> bool:
+ """Stage a stop only if no completed outcome won the row's write lock.
+
+ The caller commits the stop and any staged early-cancel flag together.
+ A winning stop merges only that flag into freshly locked extra metadata.
+ A losing stop discards its pending metadata with a locking refresh, which
+ also sees the winning row under MySQL's REPEATABLE READ isolation.
+ """
+ with db.session.no_autoflush:
+ early_cancel = query.extra.get(QUERY_EARLY_CANCEL_KEY)
+ updated = (
+ db.session.query(Query)
+ .filter(
+ Query.id == query.id,
+ ~Query.status.in_(
+ [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+ ),
+ )
+ .update(
+ {Query.status: Query.status},
+ synchronize_session=False,
+ )
+ )
+ if not updated:
+ db.session.refresh(query, with_for_update=True)
+ return False
+ # The update holds the write lock through commit. Discard stale ORM
+ # metadata and merge only stop intent into the worker's latest extra.
+ db.session.refresh(query, with_for_update=True)
+ if query.extra.get(QUERY_TARGET_COMMITTED_KEY) is True:
+ return False
+ if query.status != QueryStatus.STOPPED:
+ query.status = QueryStatus.STOPPED
+ query.end_time = end_time
+ elif query.end_time is None:
+ query.end_time = end_time
+ if early_cancel:
+ query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ return True
+
+
+def commit_query_unless_stopped(query: Query) -> bool:
+ """Atomically guard execution 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)
+ .update({Query.status: query.status}, synchronize_session=False)
+ )
+ if not updated:
+ # A locking read also sees the latest row under REPEATABLE READ.
+ db.session.refresh(query, with_for_update=True)
+ return False
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return True
+
+
+def commit_target_transaction(query: Query, commit: Callable[[], None]) -> bool:
+ """Serialize the last Stop check and target commit on the metadata row.
+
+ Postgres/MySQL hold the row lock until the target-committed marker is
+ durable. A Stop's guarded update waits, then declines that marker even
+ while result serialization is still running. SQLite ignores FOR UPDATE;
+ it retains a check-to-commit window. Neither database can make the two
+ commits atomic in the face of a metadata commit failure.
+ """
+ with db.session.no_autoflush:
+ db.session.refresh(query, attribute_names=["status"], with_for_update=True)
+ if query.status == QueryStatus.STOPPED:
+ db.session.refresh(query)
+ return False
+ commit()
+ query.set_extra_json_key(QUERY_TARGET_COMMITTED_KEY, True)
+ return commit_query_unless_stopped(query)
+
+
+def finish_query_stop(query: Query, cancel: Callable[[], bool]) -> bool:
+ """Publish Stop before cancellation; later failures cannot reject acceptance.
+
+ The caller has staged a guarded Stop. Only the first commit determines
+ acceptance. Cancellation and its optional early-cancel metadata are
+ best-effort after that point. An already-stopped row stays accepted without
+ retrying engine cancellation.
+ """
+ query_id = query.id
+ end_time = query.end_time
+ # The guarded locking refresh makes this history reflect the durable row.
+ # Reusing a stored cancel handle after Stop could cancel an unrelated session.
+ first_stop = sa.inspect(query).attrs.status.history.has_changes()
+ db.session.commit() # pylint: disable=consider-using-transaction
+ if not first_stop:
+ return True
+ try:
+ try:
+ engine_cancelled = cancel()
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Query %s: Stop accepted; engine cancellation not confirmed",
+ query_id,
+ exc_info=True,
+ )
+ else:
+ if not engine_cancelled:
+ logger.warning(
+ "Query %s: Stop accepted; engine cancellation not confirmed",
+ query_id,
+ )
+ stopped = set_query_stopped(query, end_time or now_as_float())
+ db.session.commit() # pylint: disable=consider-using-transaction
+ return stopped
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Query %s: Stop accepted; cancellation metadata was not recorded",
+ query_id,
+ exc_info=True,
+ )
+ db.session.rollback() # pylint: disable=consider-using-transaction
+ # Keep acceptance readable even if the metadata database is unavailable
+ # after the first commit. These are committed values, not staged writes.
+ sa.orm.attributes.set_committed_value(query, "status", QueryStatus.STOPPED)
+ sa.orm.attributes.set_committed_value(query, "end_time", end_time)
+ return True
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index ad64b27b7e0..98bb20bbc68 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -20,6 +20,7 @@ import logging
import sys
import traceback
import uuid
+from collections.abc import Callable
from contextlib import closing, nullcontext
from datetime import datetime
from sys import getsizeof
@@ -54,7 +55,6 @@ from superset.exceptions import (
SupersetDisallowedSQLTableException,
SupersetDMLNotAllowedException,
SupersetErrorException,
- SupersetErrorsException,
SupersetInvalidCTASException,
SupersetInvalidCVASException,
SupersetResultsBackendNotConfigureException,
@@ -63,6 +63,13 @@ 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.executor import build_statement_blocks
+from superset.sql.execution.utils import (
+ _discard_stored_results as _discard_unpublished_results,
+ _discard_unreferenced_results,
+ _extract_errors,
+ commit_query_unless_stopped,
+ commit_target_transaction,
+)
from superset.sql.parse import BaseSQLStatement, CTASMethod, SQLScript, Table
from superset.sqllab.limiting_factor import LimitingFactor
from superset.sqllab.utils import write_ipc_buffer
@@ -101,74 +108,62 @@ def handle_query_error(
query: Query,
payload: Optional[dict[str, Any]] = None,
prefix_message: str = "",
+ unpublished_results_key: Optional[str] = None,
) -> dict[str, Any]:
- """Local method handling error while processing the SQL"""
+ """
+ Local method handling error while processing the SQL
+
+ :param unpublished_results_key: key of results stored in the results
+ backend that no committed query row references; the failed row keeps
+ it, or it is deleted if a stop wins
+ """
payload = payload or {}
- # A stop request may have already committed STOPPED status while this
- # exception was being raised/propagated -- this function is the general
- # catch-all for failures anywhere in execute_sql_statements (connection
- # setup, cancel-ID acquisition, parsing, or a per-block failure), not
- # just ones caused by the stop itself. A terminal stop must stay
- # terminal, so don't let an unrelated error overwrite it with FAILED.
- #
- # Deliberately NOT a flush()-then-refresh(query) here, unlike the other
- # STOPPED-preservation checks in this module: the exception that got us
- # here may itself have already set query.status (or other attributes)
- # locally (e.g. SoftTimeLimitExceeded's own handler sets TIMED_OUT
- # without committing). Flushing first would push that stale local state
- # to the DB, clobbering a concurrently-committed STOPPED before this
- # check ever gets to observe it.
- #
- # A targeted refresh(attribute_names=["status"]) alone isn't enough:
- # verified empirically that even though it expires and reloads only the
- # named attribute (so a dirty `status` itself is correctly discarded
- # rather than written), the reload's own SELECT still triggers a normal
- # autoflush of any OTHER dirty attribute on the session first -- e.g. a
- # pending query.tmp_table_name or query.executed_sql set earlier would
- # still get written before the status read. no_autoflush suppresses
- # that: verified it emits only the targeted SELECT, with no UPDATE
- # beforehand, and leaves other pending attributes exactly as dirty as
- # they were (to be flushed normally by this function's own commit()
- # below, once we're past the STOPPED check).
with db.session.no_autoflush:
+ # The initial read avoids processing an error for an observed stop. The
+ # conditional terminal write below also covers stops committed after it.
db.session.refresh(query, attribute_names=["status"])
- if query.status == QueryStatus.STOPPED:
- payload.update({"status": query.status})
- return payload
-
- msg = f"{prefix_message} {str(ex)}".strip()
- query.error_message = msg
- query.tmp_table_name = None
- query.status = QueryStatus.FAILED
- # TODO: re-enable this after updating the frontend to properly display timeout status # noqa: E501
- # if query.status != QueryStatus.TIMED_OUT:
- # query.status = QueryStatus.FAILED
- if not query.end_time:
- query.end_time = now_as_float()
-
- # extract DB-specific errors (invalid column, eg)
- if isinstance(ex, SupersetErrorException):
- errors = [ex.error]
- elif isinstance(ex, SupersetErrorsException):
- errors = ex.errors
- else:
- errors = query.database.db_engine_spec.extract_errors(
- str(ex), database_name=query.database.unique_name
+ if query.status == QueryStatus.STOPPED:
+ db.session.refresh(query)
+ _discard_unreferenced_results(
+ query.id, unpublished_results_key, results_backend
+ )
+ payload.update({"status": QueryStatus.STOPPED})
+ return payload
+
+ msg = f"{prefix_message} {str(ex)}".strip()
+ query.error_message = msg
+ query.tmp_table_name = None
+ query.status = QueryStatus.FAILED
+ if unpublished_results_key:
+ # A rollback after a failed flush drops the pending key.
+ query.results_key = unpublished_results_key
+ # TODO: re-enable this after updating the frontend to properly display timeout status # noqa: E501
+ # if query.status != QueryStatus.TIMED_OUT:
+ # query.status = QueryStatus.FAILED
+ query.end_time = query.end_time or now_as_float()
+
+ errors = _extract_errors(ex, query)
+ errors_payload = [dataclasses.asdict(error) for error in errors]
+ if errors:
+ query.set_extra_json_key("errors", errors_payload)
+
+ payload.update(
+ {"status": QueryStatus.FAILED, "error": msg, "errors": errors_payload}
)
-
- errors_payload = [dataclasses.asdict(error) for error in errors]
- if errors:
- query.set_extra_json_key("errors", errors_payload)
-
- db.session.commit()
- payload.update({"status": query.status, "error": msg, "errors": errors_payload})
- if app.config.get("SHOW_STACKTRACE"):
- if stacktrace := traceback.format_exc():
- payload["stacktrace"] = stacktrace
- if troubleshooting_link := app.config["TROUBLESHOOTING_LINK"]:
- payload["link"] = troubleshooting_link
- return payload
+ if app.config.get("SHOW_STACKTRACE"):
+ if stacktrace := traceback.format_exc():
+ payload["stacktrace"] = stacktrace
+ if troubleshooting_link := app.config["TROUBLESHOOTING_LINK"]:
+ payload["link"] = troubleshooting_link
+ # Acquire the row lock before flushing any failure metadata.
+ if not commit_query_unless_stopped(query):
+ _discard_unreferenced_results(
+ query.id, unpublished_results_key, results_backend
+ )
+ db.session.commit()
+ return {"query_id": query.id, "status": QueryStatus.STOPPED}
+ return payload
def get_query_backoff_handler(details: dict[Any, Any]) -> None:
@@ -238,7 +233,16 @@ def get_sql_results( # pylint: disable=too-many-arguments
# request, where a nested context would replace the real Flask session with an
# empty one -- breaking RLS clauses whose Jinja macros read ``flask.session``.
with nullcontext() if has_request_context() else app.test_request_context():
- with override_user(security_manager.find_user(username)):
+ with (
+ override_user(security_manager.find_user(username)),
+ db.session.no_autoflush,
+ ):
+ unpublished_results_key: Optional[str] = None
+
+ def track_results_key(key: Optional[str]) -> None:
+ nonlocal unpublished_results_key
+ unpublished_results_key = key
+
try:
return execute_sql_statements(
query_id,
@@ -248,13 +252,26 @@ def get_sql_results( # pylint: disable=too-many-arguments
start_time=start_time,
expand_data=expand_data,
log_params=log_params,
+ on_results_key=track_results_key,
)
except Exception as ex: # pylint: disable=broad-except
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(query_id=query_id)
+ return handle_query_error(
+ ex,
+ query,
+ unpublished_results_key=unpublished_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, unpublished_results_key, results_backend
+ )
+ raise
S = TypeVar("S", bound=BaseSQLStatement[Any])
@@ -446,8 +463,14 @@ def execute_sql_statements( # noqa: C901
start_time: Optional[float],
expand_data: bool,
log_params: Optional[dict[str, Any]],
+ on_results_key: Optional[Callable[[Optional[str]], None]] = None,
) -> Optional[dict[str, Any]]:
- """Executes the sql query returns the results."""
+ """
+ Executes the sql query returns the results.
+
+ :param on_results_key: called with the key of results written to the
+ results backend, then with None once the query row commits that key
+ """
if store_results and start_time:
# only asynchronous queries
stats_logger = app.config["STATS_LOGGER"]
@@ -461,11 +484,8 @@ def execute_sql_statements( # noqa: C901
# picked it up). Honor it here, mirroring the per-block stopped-check
# further down, instead of unconditionally overwriting it back to
# RUNNING and dispatching the statement anyway.
- #
- # Same disclosed, unfixed TOCTOU residual as the other status checks in
- # this function (see the longer comment above the pre-payload check
- # further down): a stop committed strictly between this check and the
- # `query.status = RUNNING` commit a few lines below is still missed.
+ # The guarded RUNNING write below also honors a stop committed after
+ # this preliminary read.
if query.status == QueryStatus.STOPPED:
payload.update({"status": query.status})
return payload
@@ -480,7 +500,9 @@ def execute_sql_statements( # noqa: C901
logger.info("Query %s: Set query to 'running'", str(query_id))
query.status = QueryStatus.RUNNING
query.start_running_time = now_as_float()
- db.session.commit()
+ if not commit_query_unless_stopped(query):
+ db.session.commit()
+ return {"query_id": query_id, "status": QueryStatus.STOPPED}
parsed_script = SQLScript(rendered_query, engine=db_engine_spec.engine)
@@ -571,10 +593,8 @@ def execute_sql_statements( # noqa: C901
cursor = conn.cursor()
cancel_query_id = db_engine_spec.get_cancel_query_id(cursor, query)
- # Recorded unconditionally -- even when no cancel ID comes back --
- # so cancel_query() can tell "hasn't reached the engine yet" (still
- # safe to fabricate a stop) apart from "this engine has no cancel
- # support" (must fail honestly) once we get here.
+ # Record dispatch even without a cancel ID so cancel_query() can
+ # distinguish early stop intent from unavailable engine cancellation.
query.set_extra_json_key(QUERY_DISPATCHED_KEY, True)
if cancel_query_id is not None:
query.set_extra_json_key(QUERY_CANCEL_KEY, cancel_query_id)
@@ -627,46 +647,22 @@ def execute_sql_statements( # noqa: C901
payload = handle_query_error(ex, query, payload, prefix_message)
return payload
- # Commit the connection so CTA queries will create the table and any DML.
+ # A recorded Stop discards results even when an autocommit statement
+ # has already changed the target database. It cannot undo those effects.
if parsed_script.has_mutation() or query.select_as_cta:
- conn.commit()
-
- # A stop request may have landed after the last per-block check but
- # before the final statement finished (there's no next iteration to
- # catch it on for the last block). Check again before building a SUCCESS
- # payload or writing results to the backend -- both would otherwise
- # disagree with the row. The results-backend-write-failure branch below
- # has its own second check for the same reason (a stop landing while
- # that specific write is in flight).
- #
- # KNOWN, DELIBERATELY UNFIXED RESIDUAL: this codebase has no DB-level
- # locking, so every "check status, then later commit something based on
- # what was read" pattern in this function -- this one, the
- # results-backend-write-failure check below, the startup check before
- # `query.status = RUNNING` is committed a few lines later, and
- # cancel_query()'s own QUERY_DISPATCHED_KEY read/commit gap (see the
- # disclosure comment there) -- has the same fundamental TOCTOU window: a
- # stop committed strictly between the check and the later commit is
- # still missed. Each check narrows its window as much as reasonably
- # possible without locking; none of them claim to close it. Closing any
- # of them for real needs real DB-level row locking (e.g.
- # SELECT ... FOR UPDATE) or optimistic-concurrency versioning on the
- # query row, neither of which is meaningfully verifiable against the
- # sqlite backend this codebase tests against, and is deliberately not
- # attempted here.
- #
- # flush() first: refresh() does NOT autoflush -- without this, any
- # pending, uncommitted attribute set earlier in this iteration (e.g.
- # query.executed_sql, set just before execute_query() ran) would be
- # silently discarded and reloaded back to its previous committed value
- # instead of surviving to the function's own later commits.
- db.session.flush()
- db.session.refresh(query)
- if query.status == QueryStatus.STOPPED:
- payload.update({"status": query.status})
- return payload
+ if not commit_target_transaction(query, conn.commit):
+ payload.update({"status": QueryStatus.STOPPED})
+ return payload
+
+ # Preserve pending execution metadata without flushing it over a stop.
+ # Publication below acquires the row lock before any terminal ORM flush.
+ with db.session.no_autoflush:
+ db.session.refresh(query, attribute_names=["status"])
+ if query.status == QueryStatus.STOPPED:
+ db.session.refresh(query)
+ payload.update({"status": query.status})
+ return payload
- # Success, updating the query entry in database
query.rows = result_set.size
query.progress = 100
query.set_extra_json_key("progress", None)
@@ -698,6 +694,7 @@ def execute_sql_statements( # noqa: C901
)
payload["query"]["state"] = QueryStatus.SUCCESS
+ stored_key: Optional[str] = None
if store_results and results_backend:
key = str(uuid.uuid4())
payload["query"]["resultsKey"] = key
@@ -754,42 +751,14 @@ def execute_sql_statements( # noqa: C901
# For async queries (not returning results inline), mark as FAILED
# because results are inaccessible to the user
if not return_results:
- # A stop request may have landed and committed STOPPED
- # while this (potentially slow) results-backend write was
- # in flight. Refresh before marking FAILED -- a terminal
- # STOPPED must stay terminal, not be overwritten just
- # because the backend write also failed to complete
- # around the same time.
- #
- # flush() first: refresh() does NOT autoflush -- without
- # this, the result metadata already set earlier in this
- # function (rows, progress, extra "columns", select_sql,
- # end_time) plus the results_key = None set just above
- # would be silently discarded and reloaded back to their
- # previous (pre-execution) values instead of surviving to
- # this branch's own commit below.
- db.session.flush()
- db.session.refresh(query)
- if query.status == QueryStatus.STOPPED:
- # A fresh, minimal payload -- not `payload.update()`.
- # By this point `payload` already has the full
- # SUCCESS shape baked in from earlier (result data, a
- # nested query["state"] == SUCCESS, and a resultsKey
- # for a write that just failed), so patching only the
- # top-level "status" key would return a payload that
- # simultaneously claims STOPPED while still carrying
- # SUCCESS data and a resultsKey pointing at nothing
- # actually stored. Matches the shape the other
- # STOPPED-preservation return sites in this function
- # use (a plain {"query_id", "status"} pair).
- return {"query_id": query_id, "status": query.status}
-
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()
+ if not commit_query_unless_stopped(query):
+ db.session.commit()
+ return {"query_id": query_id, "status": QueryStatus.STOPPED}
raise SupersetErrorException(
SupersetError(
message=__(
@@ -802,33 +771,15 @@ def execute_sql_statements( # noqa: C901
else:
# Write succeeded - set results_key in database
query.results_key = key
+ stored_key = key
+ if on_results_key:
+ on_results_key(key)
logger.info(
"Query %s: Successfully stored results in backend, key: %s",
str(query_id),
key,
)
- # Only set SUCCESS if we didn't already set FAILED above, and don't
- # clobber a STOPPED status a concurrent stop request may have committed
- # since the check above -- a terminal stop must stay terminal. This is a
- # backstop for the DB row specifically (the payload/results-write
- # consistency check already happened above); it doesn't reopen or
- # re-narrow the same disclosed race window from that check.
- #
- # flush() first: refresh() does NOT autoflush -- without this, every
- # result field set on the success path above (rows, progress, extra
- # "columns", select_sql, end_time, results_key) would be silently
- # discarded and reloaded back to their pre-execution (typically None)
- # values on EVERY successful query, since nothing before this point
- # commits them. This was a real regression caught by CI integration
- # tests across all three DB backends (sqlite/mysql/postgres) that the
- # unit-test suite driving this fix never exercised.
- db.session.flush()
- db.session.refresh(query)
- if query.status not in (QueryStatus.FAILED, QueryStatus.STOPPED):
- query.status = QueryStatus.SUCCESS
- db.session.commit()
-
if return_results:
# since we're returning results we need to create non-arrow data
if use_arrow_data:
@@ -865,6 +816,37 @@ def execute_sql_statements( # noqa: C901
level=ErrorLevel.ERROR,
)
)
+
+ # A stop that landed during serialization or the results backend write
+ # wins unless the target-commit boundary already rejected Stop. The payload
+ # must say STOPPED too rather than carry the SUCCESS built above.
+ # Reload only the status, without flushing, so the result fields set on
+ # the success path above (rows, progress, extra "columns", select_sql,
+ # end_time, results_key) stay pending for the commit below.
+ with db.session.no_autoflush:
+ db.session.refresh(query, attribute_names=["status"])
+ if query.status == QueryStatus.STOPPED:
+ # Discard the pending success fields, keeping the stop's end time, and
+ # delete stored results the stopped row will not reference.
+ with db.session.no_autoflush:
+ db.session.refresh(query)
+ _discard_unpublished_results(query_id, stored_key, results_backend)
+ if on_results_key:
+ on_results_key(None)
+ db.session.commit()
+ return {"query_id": query_id, "status": QueryStatus.STOPPED}
+ if query.status != QueryStatus.FAILED:
+ query.status = QueryStatus.SUCCESS
+ if not commit_query_unless_stopped(query):
+ _discard_unpublished_results(query_id, stored_key, results_backend)
+ if on_results_key:
+ on_results_key(None)
+ db.session.commit()
+ return {"query_id": query_id, "status": QueryStatus.STOPPED}
+ if on_results_key:
+ on_results_key(None)
+
+ if return_results:
return payload
return None
@@ -895,39 +877,17 @@ def cancel_query(query: Query) -> bool:
cancel_query_id = query.extra.get(QUERY_CANCEL_KEY)
if cancel_query_id is None:
- # KNOWN LIMITATION (deliberately not fixed here): this read of
- # QUERY_DISPATCHED_KEY and execute_sql_statements()'s own commit of
- # that same flag (see the "Recorded unconditionally" comment where
- # it's set) are two independent transactions with no lock between
- # them. A stop request can still land in the narrow window where
- # this read has already happened -- deciding "not dispatched yet,
- # safe to fabricate a stop" -- but the worker's dispatch commit
- # lands immediately after, so the statement still gets sent to the
- # engine even though the row was just marked STOPPED. Closing this
- # for real needs DB-level row locking (e.g. SELECT ... FOR UPDATE)
- # or optimistic-concurrency versioning on the query row; neither is
- # meaningfully verifiable against the sqlite backend this codebase's
- # tests run against, so it's out of scope here rather than a
- # false claim of safety.
if query.extra.get(QUERY_DISPATCHED_KEY):
- # execute_sql_statements() already opened a connection and asked
- # this engine spec for a cancel handle, and still got nothing --
- # this engine genuinely has no way to cancel a query once it's
- # running. That's a real failure, not a race window; report it
- # honestly rather than fabricating a stop the engine can't back.
+ # No engine cancellation is available. This reports only the
+ # engine outcome; the caller's committed STOPPED remains accepted.
return False
- # No cancel handle has been recorded and execution hasn't reached the
- # engine yet, so "no ID" here can only mean "too early to have one" --
- # record the same early-cancel intent Trino's own
- # prepare_cancel_query() records for its harder case (ID only
- # obtainable after execution starts), so the stopped check at the top
- # of the statement-block loop honors the request instead of leaving
- # the query stuck at RUNNING with no avenue to ever stop it.
- #
- # Not committed here: the caller (QueryDAO.stop_query) commits this
- # together with status=STOPPED in one transaction, so another
- # request can never observe the flag set but the status still
- # RUNNING.
+ # QueryDAO commits STOPPED before calling this helper. The worker
+ # refreshes the row after its dispatch commit and before executing each
+ # block, so a stop preceding that refresh prevents dispatch. A stop
+ # after the final check can still race with execution; without an engine
+ # cancel handle, the statement can finish and its results are discarded.
+ # Stage early intent as well, preserving the guarded Trino/Impala path.
+ # The caller merges this flag into the stopped row with a guarded write.
query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
return True
diff --git a/tests/integration_tests/queries/api_tests.py b/tests/integration_tests/queries/api_tests.py
index 196c7f6657f..662c9a175f0 100644
--- a/tests/integration_tests/queries/api_tests.py
+++ b/tests/integration_tests/queries/api_tests.py
@@ -549,7 +549,6 @@ class TestQueryApi(SupersetTestCase):
data = json.loads(rv.data.decode("utf-8"))
assert data["message"] == "Query with client_id foo2 not found"
- # @mock.patch("superset.sql_lab.cancel_query")
def test_stop_query(self):
"""
Handles stop query when the DB engine spec does not
@@ -566,6 +565,7 @@ class TestQueryApi(SupersetTestCase):
select_sql="SELECT col1, col2 from table1",
executed_sql="SELECT col1, col2 from table1 LIMIT 100",
changed_on=datetime.utcnow() - timedelta(days=1),
+ status=QueryStatus.RUNNING,
)
self.login(ADMIN_USERNAME)
rv = self.client.post(
@@ -579,3 +579,30 @@ class TestQueryApi(SupersetTestCase):
assert data["result"] == "OK"
db.session.delete(query1)
db.session.commit()
+
+ def test_stop_query_already_finished(self) -> None:
+ """Return the completed query rather than reporting a successful stop."""
+ self.login(ADMIN_USERNAME)
+ query = self.insert_query(
+ get_example_database().id,
+ self.get_user("admin").id,
+ "done-stop",
+ sql="SELECT 1",
+ status=QueryStatus.SUCCESS,
+ )
+ try:
+ response = self.client.post(
+ "/api/v1/query/stop",
+ json={"client_id": query.client_id},
+ )
+
+ assert response.status_code == 422
+ assert response.json == {
+ "message": "Could not cancel query",
+ "query": query.to_dict(),
+ }
+ db.session.refresh(query)
+ assert query.status == QueryStatus.SUCCESS
+ finally:
+ db.session.delete(query)
+ db.session.commit()
diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py
index ca627c2848f..c08adb67966 100644
--- a/tests/integration_tests/sqllab_tests.py
+++ b/tests/integration_tests/sqllab_tests.py
@@ -27,6 +27,7 @@ from unittest import mock
import rison
from superset import db, security_manager
+from superset.common.db_query_status import QueryStatus
from superset.connectors.sqla.models import SqlaTable # noqa: F401
from superset.db_engine_specs import BaseEngineSpec
from superset.db_engine_specs.hive import HiveEngineSpec
@@ -660,15 +661,16 @@ class TestSqlLab(SupersetTestCase):
)
assert data["errors"][0]["error_type"] == "GENERIC_BACKEND_ERROR"
- @mock.patch("superset.sql_lab.db")
+ @mock.patch("superset.sql_lab.db.session")
@mock.patch("superset.sql_lab.get_query")
@mock.patch("superset.sql_lab.execute_query")
def test_execute_sql_statements(
self,
- mock_execute_query,
- mock_get_query,
- mock_db,
- ):
+ mock_execute_query: mock.MagicMock,
+ mock_get_query: mock.MagicMock,
+ mock_session: mock.MagicMock,
+ ) -> None:
+ """Split statements without using a real session in shared helpers."""
sql = dedent(
"""
-- comment
@@ -676,7 +678,6 @@ class TestSqlLab(SupersetTestCase):
SELECT /*+ hint */ @value AS foo;
"""
)
- mock_db = mock.MagicMock() # noqa: F841
mock_query = mock.MagicMock(select_as_cta=False)
mock_query.database.allow_run_async = False
mock_cursor = mock.MagicMock()
@@ -698,6 +699,8 @@ class TestSqlLab(SupersetTestCase):
expand_data=False,
log_params=None,
)
+ mock_session.query.assert_called_with(Query)
+ assert mock_query.status == QueryStatus.SUCCESS
mock_execute_query.assert_has_calls(
[
mock.call(mock_query, mock_cursor, None),
@@ -757,15 +760,16 @@ class TestSqlLab(SupersetTestCase):
},
)
- @mock.patch("superset.sql_lab.db")
+ @mock.patch("superset.sql_lab.db.session")
@mock.patch("superset.sql_lab.get_query")
@mock.patch("superset.sql_lab.execute_query")
def test_execute_sql_statements_ctas(
self,
- mock_execute_query,
- mock_get_query,
- mock_db,
- ):
+ mock_execute_query: mock.MagicMock,
+ mock_get_query: mock.MagicMock,
+ mock_session: mock.MagicMock,
+ ) -> None:
+ """Split statements without using a real session in shared helpers."""
sql = dedent(
"""
-- comment
@@ -773,7 +777,6 @@ class TestSqlLab(SupersetTestCase):
SELECT /*+ hint */ @value AS foo;
"""
)
- mock_db = mock.MagicMock() # noqa: F841
mock_query = mock.MagicMock(
select_as_cta=True,
ctas_method=CTASMethod.TABLE.name,
@@ -805,6 +808,8 @@ class TestSqlLab(SupersetTestCase):
expand_data=False,
log_params=None,
)
+ mock_session.query.assert_called_with(Query)
+ assert mock_query.status == QueryStatus.SUCCESS
mock_execute_query.assert_has_calls(
[
mock.call(mock_query, mock_cursor, None),
diff --git a/tests/testcontainers/db_engine_specs/test_mysql.py b/tests/testcontainers/db_engine_specs/test_mysql.py
index e3422660ee3..85faab114f9 100644
--- a/tests/testcontainers/db_engine_specs/test_mysql.py
+++ b/tests/testcontainers/db_engine_specs/test_mysql.py
@@ -245,3 +245,116 @@ def test_require_mysql_tls_connects_with_verified_tls(
cipher = conn.execute(text("SHOW STATUS LIKE 'Ssl_cipher'")).fetchone()
assert cipher is not None
assert cipher[1], "connection succeeded but is not actually using TLS"
+
+
+@pytest.mark.parametrize("isolation_level", ["REPEATABLE READ", "READ COMMITTED"])
+@pytest.mark.parametrize("winner", ["stop", "worker"])
+def test_stop_at_target_commit_is_serialized(
+ mysql_container: MySqlContainer,
+ monkeypatch: pytest.MonkeyPatch,
+ winner: str,
+ isolation_level: str,
+) -> None:
+ """MySQL locking reads serialize Stop with the target commit under RR and RC."""
+ from concurrent.futures import ThreadPoolExecutor, TimeoutError
+ from threading import Event
+
+ from sqlalchemy.orm import scoped_session, sessionmaker
+
+ from superset.constants import QUERY_TARGET_COMMITTED_KEY
+ from superset.extensions import db
+ from superset.models.sql_lab import Query
+ from superset.sql.execution.utils import (
+ commit_target_transaction,
+ set_query_stopped,
+ )
+
+ engine = create_engine(
+ URL.create(
+ "mysql+pymysql",
+ username=mysql_container.username,
+ password=mysql_container.password,
+ host=_tcp_host(mysql_container),
+ port=int(mysql_container.get_exposed_port(mysql_container.port)),
+ database=mysql_container.dbname,
+ )
+ )
+ with engine.begin() as connection:
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS dbs (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS ab_user (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(text("INSERT IGNORE INTO dbs VALUES (1)"))
+ Query.__table__.create(engine, checkfirst=True)
+ sessions = scoped_session(
+ sessionmaker(bind=engine.execution_options(isolation_level=isolation_level))
+ )
+ monkeypatch.setattr(db, "session", sessions)
+ sessions.add(
+ Query(
+ id=126,
+ client_id="commit-stop",
+ database_id=1,
+ status="running",
+ sql="SELECT 1",
+ )
+ )
+ sessions.commit()
+ committing = Event()
+ release = Event()
+ stopping = Event()
+
+ def target_commit() -> None:
+ """Keep the target commit pending while Stop attempts its guarded write."""
+ committing.set()
+ assert release.wait(10)
+
+ def work() -> bool:
+ """Hold the metadata lock across the target commit."""
+ try:
+ return commit_target_transaction(sessions.get(Query, 126), target_commit)
+ finally:
+ sessions.remove()
+
+ def stop() -> bool:
+ """Read the row before the worker commits its target-committed marker."""
+ try:
+ query = sessions.get(Query, 126)
+ stopping.set()
+ accepted = set_query_stopped(query, 1234.5)
+ sessions.commit()
+ return accepted
+ finally:
+ sessions.remove()
+
+ try:
+ if winner == "stop":
+ assert stop() is True
+ assert work() is False
+ assert not committing.is_set()
+ else:
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ worker = pool.submit(work)
+ try:
+ assert committing.wait(10)
+ stopper = pool.submit(stop)
+ assert stopping.wait(10)
+ with pytest.raises(TimeoutError):
+ stopper.result(timeout=0.2)
+ finally:
+ release.set()
+ assert worker.result(timeout=10) is True
+ assert stopper.result(timeout=10) is False
+ sessions.expire_all()
+ query = sessions.get(Query, 126)
+ assert query.status == ("stopped" if winner == "stop" else "running")
+ assert bool(query.extra.get(QUERY_TARGET_COMMITTED_KEY)) is (winner == "worker")
+ finally:
+ release.set()
+ sessions.rollback()
+ sessions.query(Query).filter_by(id=126).delete()
+ sessions.commit()
+ sessions.remove()
+ engine.dispose()
diff --git a/tests/testcontainers/db_engine_specs/test_postgres.py b/tests/testcontainers/db_engine_specs/test_postgres.py
index 40741815775..19fb0001913 100644
--- a/tests/testcontainers/db_engine_specs/test_postgres.py
+++ b/tests/testcontainers/db_engine_specs/test_postgres.py
@@ -97,3 +97,693 @@ def test_get_columns_maps_native_types(engine: Engine) -> None:
assert spec is not None
assert spec.generic_type == GenericDataType.NUMERIC
assert isinstance(spec.sqla_type, Integer)
+
+
+@pytest.mark.parametrize("return_results", [False, True], ids=["async", "sync"])
+@pytest.mark.parametrize("terminate", [False, True], ids=["cancel", "terminate"])
+@pytest.mark.parametrize("stop_timing", ["cursor", "error_handler"])
+def test_legacy_cancel_error_preserves_committed_stop(
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ return_results: bool,
+ terminate: bool,
+ stop_timing: str,
+) -> None:
+ """A real cancelled Postgres cursor cannot publish FAILED over a Stop."""
+ import time
+ from concurrent.futures import ThreadPoolExecutor
+ from contextlib import contextmanager
+ from typing import Any
+ from unittest.mock import MagicMock, PropertyMock
+
+ from psycopg2 import OperationalError
+ from sqlalchemy import text
+ from sqlalchemy.orm import Session
+
+ from superset import sql_lab
+ from superset.models.sql_lab import Query
+ from tests.integration_tests.test_app import app
+
+ # Minimal FK targets keep this reproduction independent of app migrations.
+ with engine.begin() as connection:
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS dbs (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS ab_user (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(text("INSERT INTO dbs VALUES (1) ON CONFLICT DO NOTHING"))
+ Query.__table__.create(engine, checkfirst=True)
+
+ def stop(*args: object, **kwargs: object) -> list[object]:
+ """Commit after the legacy handler has read RUNNING, before its write."""
+ with Session(engine) as stop_session:
+ stopped = stop_session.get(Query, 123)
+ assert stopped is not None
+ stopped.status = "stopped"
+ stopped.end_time = 1234.5
+ stop_session.commit()
+ return []
+
+ cursor_cancelled = False
+
+ def execute(cursor: Any, sql: str, query: Query) -> None:
+ """Cancel a statement observed running on an independent connection."""
+ nonlocal cursor_cancelled
+ pid = cursor.connection.get_backend_pid()
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ running = pool.submit(cursor.execute, sql)
+ with engine.connect().execution_options(
+ isolation_level="AUTOCOMMIT"
+ ) as canceller:
+ deadline = time.monotonic() + 10
+ while time.monotonic() < deadline:
+ if canceller.execute(
+ text(
+ "SELECT wait_event = 'PgSleep' FROM pg_stat_activity "
+ "WHERE pid = :pid"
+ ),
+ {"pid": pid},
+ ).scalar():
+ break
+ time.sleep(0.01)
+ else:
+ pytest.fail("Postgres statement did not start")
+ function = "pg_terminate_backend" if terminate else "pg_cancel_backend"
+ assert canceller.execute(
+ text(f"SELECT {function}(:pid)"), {"pid": pid}
+ ).scalar()
+ with pytest.raises(OperationalError) as error:
+ running.result(timeout=10)
+ cursor_cancelled = True
+ raise error.value
+
+ @contextmanager
+ def raw_connection(**kwargs: Any) -> Iterator[Any]:
+ """Use a real DBAPI connection for the complete legacy execution path."""
+ connection = engine.raw_connection()
+ try:
+ yield connection
+ finally:
+ connection.close()
+
+ database = MagicMock()
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "public"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.get_raw_connection = raw_connection
+ monkeypatch.setattr(Query, "database", PropertyMock(return_value=database))
+ monkeypatch.setattr(PostgresEngineSpec, "execute_with_cursor", execute)
+ monkeypatch.setattr(sql_lab, "warm_and_release_connection", lambda *args: None)
+ monkeypatch.setattr(
+ sql_lab,
+ "_extract_errors",
+ lambda *args, **kwargs: stop() if stop_timing == "error_handler" else [],
+ )
+ monkeypatch.setattr(sql_lab.security_manager, "find_user", lambda _: None)
+ with app.test_request_context(), Session(engine) as worker_session:
+ monkeypatch.setattr(worker_session, "remove", lambda: None, raising=False)
+ monkeypatch.setattr(sql_lab.db, "session", worker_session)
+ refresh = worker_session.refresh
+
+ def refresh_then_stop(*args: Any, **kwargs: Any) -> None:
+ """Commit after the cancelled cursor's handler reads its status."""
+ refresh(*args, **kwargs)
+ if (
+ stop_timing == "cursor"
+ and cursor_cancelled
+ and kwargs.get("attribute_names") == ["status"]
+ ):
+ stop()
+
+ monkeypatch.setattr(worker_session, "refresh", refresh_then_stop)
+ monkeypatch.setitem(app.config, "STATS_LOGGER", MagicMock())
+ monkeypatch.setitem(app.config, "DISALLOWED_SQL_FUNCTIONS", {})
+ worker_session.add(
+ Query(
+ id=123,
+ client_id="pg-cancel",
+ database_id=1,
+ sql="SELECT pg_sleep(30)",
+ status="pending",
+ )
+ )
+ worker_session.commit()
+ try:
+ payload = sql_lab.get_sql_results(
+ 123, "SELECT pg_sleep(30)", return_results
+ )
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == "stopped"
+ assert persisted.end_time == 1234.5
+ assert persisted.error_message is None
+ assert payload is not None
+ assert payload["status"] == "stopped"
+ assert "error" not in payload
+ finally:
+ worker_session.rollback()
+ worker_session.query(Query).filter_by(id=123).delete()
+ worker_session.commit()
+
+
+@pytest.mark.parametrize("delay_ms", [0, 20, 100])
+@pytest.mark.parametrize("terminate", [False, True], ids=["cancel", "terminate"])
+@pytest.mark.parametrize("caller", ["dao", "async_handle"])
+def test_stop_request_precedes_postgres_cancellation_failure(
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ delay_ms: int,
+ terminate: bool,
+ caller: str,
+) -> None:
+ """The production Stop entry points must publish intent before cancellation."""
+ import time
+ from concurrent.futures import ThreadPoolExecutor
+ from contextlib import contextmanager
+ from threading import Event
+ from typing import Any
+ from unittest.mock import MagicMock, PropertyMock
+
+ from sqlalchemy import text
+ from sqlalchemy.orm import scoped_session, Session, sessionmaker
+
+ from superset import sql_lab
+ from superset.daos.query import QueryDAO
+ from superset.models.sql_lab import Query
+ from superset.sql.execution.executor import SQLExecutor
+ from tests.integration_tests.test_app import app
+
+ with engine.begin() as connection:
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS dbs (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS ab_user (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(text("INSERT INTO dbs VALUES (1) ON CONFLICT DO NOTHING"))
+ connection.execute(
+ text("INSERT INTO ab_user VALUES (7) ON CONFLICT DO NOTHING")
+ )
+ Query.__table__.create(engine, checkfirst=True)
+ sessions = scoped_session(sessionmaker(bind=engine))
+ executing = Event()
+ finished = Event()
+ backend_pid: list[int] = []
+
+ @contextmanager
+ def raw_connection(**kwargs: Any) -> Iterator[Any]:
+ """Open a real target connection, distinct from metadata sessions."""
+ connection = engine.raw_connection()
+ try:
+ yield connection
+ finally:
+ connection.close()
+
+ def execute(cursor: Any, sql: str, query: Query) -> None:
+ """Expose the executing backend to the independent Stop request."""
+ backend_pid.append(cursor.connection.get_backend_pid())
+ executing.set()
+ cursor.execute(sql)
+
+ def cancel(cursor: Any, query: Query, cancel_query_id: str) -> bool:
+ """Let the worker finish before the Stop request resumes after cancel."""
+ function = "pg_terminate_backend" if terminate else "pg_cancel_backend"
+ cursor.execute(f"SELECT {function}(%s)", (int(cancel_query_id),))
+ assert cursor.fetchone()[0]
+ time.sleep(delay_ms / 1000)
+ # Include the schedule where even a zero-delay request is descheduled
+ # after cancel. This makes the race reproducible, not scheduler-dependent.
+ assert finished.wait(10)
+ return True
+
+ database = MagicMock()
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "public"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.get_raw_connection = raw_connection
+ database.get_sqla_engine.return_value.__enter__.return_value = engine
+ monkeypatch.setattr(Query, "database", PropertyMock(return_value=database))
+ monkeypatch.setattr(PostgresEngineSpec, "execute_with_cursor", execute)
+ monkeypatch.setattr(PostgresEngineSpec, "cancel_query", cancel)
+ monkeypatch.setattr(sql_lab, "warm_and_release_connection", lambda *args: None)
+ monkeypatch.setattr(sql_lab.security_manager, "find_user", lambda _: None)
+ monkeypatch.setattr(sql_lab.db, "session", sessions)
+ monkeypatch.setattr("superset.daos.query.get_user_id", lambda: 7)
+ monkeypatch.setitem(app.config, "STATS_LOGGER", MagicMock())
+ monkeypatch.setitem(app.config, "DISALLOWED_SQL_FUNCTIONS", {})
+
+ def work() -> dict[str, Any] | None:
+ """Run the legacy worker with a thread-local metadata session."""
+ try:
+ with app.test_request_context():
+ return sql_lab.get_sql_results(124, "SELECT pg_sleep(30)", True)
+ finally:
+ sessions.remove()
+ finished.set()
+
+ with app.test_request_context():
+ sessions.add(
+ Query(
+ id=124,
+ client_id="prod-stop",
+ user_id=7,
+ database_id=1,
+ sql="SELECT pg_sleep(30)",
+ status="pending",
+ )
+ )
+ sessions.commit()
+ try:
+ with ThreadPoolExecutor(max_workers=1) as pool:
+ worker = pool.submit(work)
+ assert executing.wait(10)
+ with engine.connect().execution_options(
+ isolation_level="AUTOCOMMIT"
+ ) as observer:
+ deadline = time.monotonic() + 10
+ while time.monotonic() < deadline:
+ if observer.execute(
+ text(
+ "SELECT wait_event = 'PgSleep' FROM pg_stat_activity "
+ "WHERE pid = :pid"
+ ),
+ {"pid": backend_pid[0]},
+ ).scalar():
+ break
+ time.sleep(0.01)
+ else:
+ pytest.fail("Postgres statement did not start")
+ if caller == "dao":
+ stopped = QueryDAO.stop_query("prod-stop")
+ assert stopped.status == "stopped"
+ else:
+ assert SQLExecutor._cancel_async_query(124, database)
+ payload = worker.result(timeout=10)
+ with Session(engine) as observer_session:
+ persisted = observer_session.get(Query, 124)
+ assert persisted is not None
+ assert persisted.status == "stopped"
+ assert persisted.error_message is None
+ assert persisted.end_time is not None
+ assert payload is not None
+ assert payload["status"] == "stopped"
+ finally:
+ sessions.rollback()
+ sessions.query(Query).filter_by(id=124).delete()
+ sessions.commit()
+ sessions.remove()
+
+
+@pytest.mark.parametrize(
+ ("statement", "autocommit"),
+ [
+ ("select", False),
+ ("insert", False),
+ ("ctas", False),
+ ("insert", True),
+ ("ctas", True),
+ ],
+ ids=[
+ "select",
+ "transaction-insert",
+ "transaction-ctas",
+ "autocommit-insert",
+ "autocommit-ctas",
+ ],
+)
+@pytest.mark.parametrize("worker_path", ["legacy", "celery"])
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("cancel_mode", ["normal", "unreachable", "no_cancel_id"])
+def test_accepted_stop_with_real_postgres_workers( # noqa: C901
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ worker_path: str,
+ caller: str,
+ cancel_mode: str,
+ statement: str,
+ autocommit: bool,
+ stop_timing: str = "executing", # noqa: PT028
+ repeat_stop: bool = False, # noqa: PT028
+) -> None:
+ """Real workers discard results after an accepted, best-effort engine stop."""
+ import inspect as python_inspect
+ import socket
+ import time
+ from concurrent.futures import ThreadPoolExecutor
+ from contextlib import contextmanager
+ from typing import Any
+ from unittest.mock import MagicMock, PropertyMock
+
+ from sqlalchemy import text
+ from sqlalchemy.orm import scoped_session, Session, sessionmaker
+
+ from superset import sql_lab
+ from superset.models.sql_lab import Query
+ from superset.queries.api import QueryRestApi
+ from superset.sql.execution import celery_task
+ from superset.sql.execution.executor import SQLExecutor
+ from tests.integration_tests.test_app import app
+
+ with engine.begin() as connection:
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS dbs (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(
+ text("CREATE TABLE IF NOT EXISTS ab_user (id INTEGER PRIMARY KEY)")
+ )
+ connection.execute(text("INSERT INTO dbs VALUES (1) ON CONFLICT DO NOTHING"))
+ connection.execute(
+ text("INSERT INTO ab_user VALUES (7) ON CONFLICT DO NOTHING")
+ )
+ Query.__table__.create(engine, checkfirst=True)
+ sessions = scoped_session(sessionmaker(bind=engine))
+ table_name = "accepted_stop_mutation"
+ with engine.begin() as connection:
+ connection.execute(text("DROP TABLE IF EXISTS accepted_stop_mutation"))
+ if statement == "insert":
+ connection.execute(
+ text("CREATE TABLE accepted_stop_mutation (value INTEGER)")
+ )
+ sql = {
+ "select": "SELECT pg_sleep(3) AS accepted_stop",
+ "insert": "INSERT INTO accepted_stop_mutation SELECT 1 FROM pg_sleep(3)",
+ "ctas": (
+ "CREATE TABLE accepted_stop_mutation AS SELECT 1 AS value FROM pg_sleep(3)"
+ ),
+ }[statement]
+
+ from threading import Event
+
+ committing = Event()
+ release_commit = Event()
+ release_results = Event()
+
+ class CursorProxy:
+ """Inject a Celery signal after the real statement has returned."""
+
+ def __init__(self, cursor: Any) -> None:
+ self.cursor = cursor
+
+ def __getattr__(self, name: str) -> Any:
+ return getattr(self.cursor, name)
+
+ def fetchall(self) -> Any:
+ if stop_timing == "soft_timeout":
+ from celery.exceptions import SoftTimeLimitExceeded
+
+ raise SoftTimeLimitExceeded()
+ return self.cursor.fetchall()
+
+ class ConnectionProxy:
+ """Pause exactly between the worker's stop check and target commit."""
+
+ def __init__(self, connection: Any) -> None:
+ self.connection = connection
+
+ def cursor(self) -> Any:
+ return CursorProxy(self.connection.cursor())
+
+ def commit(self) -> None:
+ if stop_timing == "target_commit":
+ committing.set()
+ assert release_commit.wait(10), "Commit was not released"
+ self.connection.commit()
+
+ @contextmanager
+ def raw_connection(**kwargs: Any) -> Iterator[Any]:
+ """Use a real Postgres cursor without replacing either worker's executor."""
+ connection = engine.raw_connection()
+ original_autocommit = connection.dbapi_connection.autocommit
+ connection.dbapi_connection.autocommit = autocommit
+ try:
+ yield ConnectionProxy(connection)
+ finally:
+ # Direct DBAPI changes bypass SQLAlchemy's isolation-level reset.
+ # Do not return an autocommit connection to the metadata pool.
+ connection.rollback()
+ connection.dbapi_connection.autocommit = original_autocommit
+ connection.close()
+
+ database = MagicMock()
+ database.database_name = "Postgres stop regression"
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ database.allow_dml = True
+ database.get_default_schema_for_query.return_value = "public"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.get_raw_connection = raw_connection
+ database.get_sqla_engine.return_value.__enter__.return_value = engine
+ monkeypatch.setattr(Query, "database", PropertyMock(return_value=database))
+ monkeypatch.setattr(Query, "user", PropertyMock(return_value=None))
+ monkeypatch.setattr(sql_lab, "warm_and_release_connection", lambda *args: None)
+ monkeypatch.setattr(sql_lab.security_manager, "find_user", lambda _: None)
+ monkeypatch.setattr(sql_lab.db, "session", sessions)
+ monkeypatch.setattr("superset.daos.query.get_user_id", lambda: 7)
+ monkeypatch.setitem(app.config, "STATS_LOGGER", MagicMock())
+ monkeypatch.setitem(app.config, "DISALLOWED_SQL_FUNCTIONS", {})
+ backend = MagicMock()
+ monkeypatch.setattr(celery_task, "results_backend", backend)
+ monkeypatch.setattr(sql_lab, "results_backend", backend)
+
+ def store_results(*args: Any) -> bool:
+ """Keep finalization pending until Stop has checked the commit marker."""
+ if stop_timing == "target_commit":
+ assert release_results.wait(10), "Results were not released"
+ return True
+
+ backend.set.side_effect = store_results
+ cancel = MagicMock(wraps=PostgresEngineSpec.cancel_query)
+ if repeat_stop:
+ monkeypatch.setattr(PostgresEngineSpec, "cancel_query", cancel)
+ if cancel_mode == "no_cancel_id":
+ # Model a connector using BaseEngineSpec's default (no cancel handle).
+ monkeypatch.setattr(
+ PostgresEngineSpec, "get_cancel_query_id", lambda *args: None
+ )
+
+ def work() -> dict[str, Any] | None:
+ """Run the production Celery task (including _execute_sql_statements)."""
+ try:
+ with app.test_request_context():
+ if worker_path == "celery":
+ return celery_task.execute_sql_task(125, sql)
+ return sql_lab.get_sql_results(125, sql, True, store_results=True)
+ finally:
+ sessions.remove()
+
+ # Reserve a local port without listening: cancel gets a real connection
+ # refusal while the target statement continues on its existing connection.
+ with socket.socket() as unavailable, app.test_request_context():
+ unavailable.bind(("127.0.0.1", 0))
+ unreachable = create_engine(
+ engine.url.set(host="127.0.0.1", port=unavailable.getsockname()[1]),
+ connect_args={"connect_timeout": 1},
+ )
+ if cancel_mode == "unreachable":
+ database.get_sqla_engine.return_value.__enter__.return_value = unreachable
+ sessions.add(
+ Query(
+ id=125,
+ client_id="accept-stop",
+ user_id=7,
+ database_id=1,
+ sql=sql,
+ status="pending",
+ )
+ )
+ sessions.commit()
+ try:
+ with ThreadPoolExecutor(max_workers=2) as pool:
+ worker = pool.submit(work)
+ with engine.connect().execution_options(
+ isolation_level="AUTOCOMMIT"
+ ) as observer:
+ deadline = time.monotonic() + 10
+ pid = None
+ while (
+ stop_timing != "target_commit" and time.monotonic() < deadline
+ ):
+ pid = observer.execute(
+ text(
+ "SELECT pid FROM pg_stat_activity "
+ "WHERE wait_event = 'PgSleep' "
+ "AND query LIKE '%accepted_stop%'"
+ )
+ ).scalar()
+ if pid:
+ break
+ time.sleep(0.01)
+ if stop_timing == "target_commit":
+ assert committing.wait(10), "Worker did not reach target commit"
+ else:
+ assert pid is not None, "Postgres statement did not start"
+
+ def stop_request() -> bool:
+ """Run the real API/handle Stop in its own metadata session."""
+ try:
+ with app.test_request_context(
+ json={"client_id": "accept-stop"}
+ ):
+ if caller == "api":
+ response = python_inspect.unwrap(
+ QueryRestApi.stop_query
+ )(QueryRestApi())
+ assert response.status_code in (200, 422)
+ if (
+ stop_timing == "target_commit"
+ and response.status_code == 422
+ ):
+ assert (
+ response.get_json()["query"]["state"]
+ == "running"
+ )
+ return response.status_code == 200
+ return (
+ SQLExecutor(database)
+ ._create_async_handle(125)
+ .cancel()
+ )
+ finally:
+ sessions.remove()
+
+ stopping = pool.submit(stop_request)
+ if stop_timing == "target_commit":
+ # A guarded Stop must wait for the worker's row lock.
+ # On the old code it completes before the target commit.
+ deadline = time.monotonic() + 5
+ while not stopping.done() and time.monotonic() < deadline:
+ if observer.execute(
+ text(
+ "SELECT count(*) FROM pg_stat_activity WHERE "
+ "wait_event_type = 'Lock' AND query LIKE '%query%'"
+ )
+ ).scalar():
+ break
+ time.sleep(0.01)
+ release_commit.set()
+ try:
+ accepted = stopping.result(timeout=10)
+ assert accepted is (stop_timing != "target_commit")
+ finally:
+ release_results.set()
+ payload = worker.result(timeout=10)
+ assert not observer.execute(
+ text(
+ "SELECT count(*) FROM pg_stat_activity "
+ "WHERE pid = :pid AND state = 'active'"
+ ),
+ {"pid": pid},
+ ).scalar()
+ with Session(engine) as verifier:
+ stopped = verifier.get(Query, 125)
+ assert stopped is not None
+ expected = "success" if stop_timing == "target_commit" else "stopped"
+ assert stopped.status == expected
+ assert stopped.to_dict()["state"] == expected
+ assert stopped.end_time is not None
+ if repeat_stop:
+ original_end_time = stopped.end_time
+ assert stop_request() is True
+ verifier.refresh(stopped)
+ assert stopped.end_time == original_end_time
+ cancel.assert_called_once()
+ assert stopped.error_message is None
+ assert bool(stopped.results_key) is (stop_timing == "target_commit")
+ assert payload is not None
+ assert payload["status"] == expected
+ assert "error" not in payload
+ if stop_timing != "target_commit":
+ backend.set.assert_not_called()
+ # Stop accepts discarding results, not rolling back an autocommit
+ # mutation that engine cancellation could not interrupt.
+ durable = stop_timing == "target_commit" or (
+ autocommit and cancel_mode != "normal"
+ )
+ with engine.connect() as observer:
+ if statement == "ctas":
+ assert inspect(observer).has_table(table_name) is durable
+ elif statement == "insert":
+ assert observer.execute(
+ text("SELECT COUNT(*) FROM accepted_stop_mutation")
+ ).scalar() == int(durable)
+ finally:
+ sessions.rollback()
+ sessions.query(Query).filter_by(id=125).delete()
+ sessions.commit()
+ sessions.remove()
+ unreachable.dispose()
+ with engine.begin() as connection:
+ connection.execute(text("DROP TABLE IF EXISTS accepted_stop_mutation"))
+
+
+@pytest.mark.parametrize("worker_path", ["legacy", "celery"])
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("cancel_mode", ["normal", "unreachable", "no_cancel_id"])
+@pytest.mark.parametrize("statement", ["insert", "ctas"])
+def test_stop_at_target_commit_is_serialized(
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ worker_path: str,
+ caller: str,
+ cancel_mode: str,
+ statement: str,
+) -> None:
+ """A target commit holding the row lock wins over a concurrent Stop."""
+ test_accepted_stop_with_real_postgres_workers(
+ engine,
+ monkeypatch,
+ worker_path,
+ caller,
+ cancel_mode,
+ statement,
+ False,
+ stop_timing="target_commit",
+ )
+
+
+@pytest.mark.parametrize("worker_path", ["legacy", "celery"])
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("cancel_mode", ["unreachable", "no_cancel_id"])
+def test_soft_timeout_after_accepted_postgres_stop(
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ worker_path: str,
+ caller: str,
+ cancel_mode: str,
+) -> None:
+ """A time-limit signal cannot overwrite an already accepted Stop."""
+ test_accepted_stop_with_real_postgres_workers(
+ engine,
+ monkeypatch,
+ worker_path,
+ caller,
+ cancel_mode,
+ "select",
+ False,
+ stop_timing="soft_timeout",
+ )
+
+
+@pytest.mark.parametrize("worker_path", ["legacy", "celery"])
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+def test_repeated_stop_cancels_postgres_only_once(
+ engine: Engine,
+ monkeypatch: pytest.MonkeyPatch,
+ worker_path: str,
+ caller: str,
+) -> None:
+ """An accepted Stop cannot reuse a closed worker connection's cancel PID."""
+ test_accepted_stop_with_real_postgres_workers(
+ engine,
+ monkeypatch,
+ worker_path=worker_path,
+ caller=caller,
+ cancel_mode="normal",
+ statement="select",
+ autocommit=False,
+ repeat_stop=True,
+ )
diff --git a/tests/unit_tests/dao/queries_test.py b/tests/unit_tests/dao/queries_test.py
index 40a2f0b0dc8..ba44312edf0 100644
--- a/tests/unit_tests/dao/queries_test.py
+++ b/tests/unit_tests/dao/queries_test.py
@@ -21,7 +21,7 @@ import pytest
from pytest_mock import MockerFixture
from sqlalchemy.orm.session import Session
-from superset.exceptions import QueryNotFoundException, SupersetCancelQueryException
+from superset.exceptions import QueryNotFoundException
def test_query_dao_save_metadata(session: Session) -> None:
@@ -146,10 +146,7 @@ def test_query_dao_stop_query_not_found(
db.session.add(database)
db.session.add(query_obj)
- # Committed (not just autoflushed) since QueryDAO.stop_query() is now
- # wrapped in @transaction, which rolls back the session on the
- # QueryNotFoundException raised below -- an uncommitted insert would be
- # discarded along with it.
+ # Persist the fixture before testing a different query identifier.
db.session.commit()
mocker.patch("superset.sql_lab.cancel_query", return_value=False)
@@ -202,9 +199,10 @@ def test_query_dao_stop_query_not_running(
assert query.status == QueryStatus.FAILED
-def test_query_dao_stop_query_failed(
+def test_query_dao_stop_query_accepts_uncancelable_engine(
mocker: MockerFixture, app: Any, session: Session
) -> None:
+ """A missing engine cancellation does not reject a durable Stop."""
from superset import db
from superset.common.db_query_status import QueryStatus
from superset.models.core import Database
@@ -233,21 +231,20 @@ def test_query_dao_stop_query_failed(
db.session.add(database)
db.session.add(query_obj)
- # Committed (not just autoflushed) since QueryDAO.stop_query() is now
- # wrapped in @transaction, which rolls back the session on the
- # SupersetCancelQueryException raised below -- an uncommitted insert
- # would be discarded along with it.
+ # The query must exist before the Stop transaction can lock its row.
db.session.commit()
- mocker.patch("superset.sql_lab.cancel_query", return_value=False)
+ cancel_query = mocker.patch("superset.sql_lab.cancel_query", return_value=False)
from superset.daos.query import QueryDAO
- with pytest.raises(SupersetCancelQueryException):
- QueryDAO.stop_query(query_obj.client_id)
+ assert QueryDAO.stop_query(query_obj.client_id) is query_obj
+ cancel_query.assert_called_once_with(query_obj)
+ db.session.expire_all()
query = db.session.query(Query).one()
- assert query.status == QueryStatus.RUNNING
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time is not None
def test_query_dao_stop_query(
@@ -324,10 +321,7 @@ def test_query_dao_stop_query_wrong_user(
db.session.add(database)
db.session.add(query_obj)
- # Committed (not just autoflushed) since QueryDAO.stop_query() is now
- # wrapped in @transaction, which rolls back the session on the
- # QueryNotFoundException raised below -- an uncommitted insert would be
- # discarded along with it.
+ # Persist the fixture before testing ownership scoping.
db.session.commit()
# Simulate a different user (user 2) attempting to stop user 1's query
diff --git a/tests/unit_tests/dao/test_query_stop_race.py b/tests/unit_tests/dao/test_query_stop_race.py
index 201753facc5..324fc1d7def 100644
--- a/tests/unit_tests/dao/test_query_stop_race.py
+++ b/tests/unit_tests/dao/test_query_stop_race.py
@@ -22,23 +22,12 @@ query.extra -- superset.sql_lab.cancel_query() would return False without
recording anything, so stop_query() raised SupersetCancelQueryException and
left query.status untouched.
-The fix distinguishes two cases that both look like "no cancel_query_id yet":
-
-1. Execution hasn't reached the engine at all (hasn't opened a DB connection
- or called get_cancel_query_id()). Here, generalizing the
- QUERY_EARLY_CANCEL_KEY mechanism Trino's own prepare_cancel_query() already
- uses for its harder case is safe: record the early-cancel intent and
- report success, so the stopped-check at the top of the statement-block
- loop honors it before the statement is ever sent to the database.
-
-2. Execution HAS reached the engine (a connection is open and
- get_cancel_query_id() has been called) and the engine spec still returned
- no cancel ID -- this engine genuinely has no way to cancel a running
- query. This must continue to fail honestly (raise
- SupersetCancelQueryException, leave status alone) rather than fabricate a
- stop the engine can't actually back; case 1's success path must not apply
- here. A new QUERY_DISPATCHED_KEY, set unconditionally once execution opens
- the connection, is what distinguishes the two cases.
+Stop intent is durable independently of engine cancellation support. Before
+execution reaches the engine, the early-cancel flag prevents dispatch. After
+execution has started, an engine without a cancel handle may finish its SQL,
+but the accepted Stop still discards its results and retains STOPPED. The
+QUERY_DISPATCHED_KEY distinguishes missing cancellation support from the
+pre-dispatch window; neither case may reject an already-durable Stop.
Also covered: the success-path status finalizer at the end of
execute_sql_statements() must not clobber a STOPPED status a concurrent stop
@@ -71,7 +60,6 @@ superset/sql_lab.py::cancel_query/execute_sql_statements/handle_query_error.
import threading
from typing import Any
-import pytest
from pytest_mock import MockerFixture
from sqlalchemy.orm.session import Session
@@ -315,27 +303,20 @@ def test_stop_during_pre_dispatch_pause_marks_stopped_without_dispatch(
assert execution_result["payload"]["status"] == QueryStatus.STOPPED
-def test_stop_after_dispatch_with_no_cancel_support_raises_honestly(
+def test_stop_after_dispatch_with_no_cancel_support_discards_results(
mocker: MockerFixture, app: Any
) -> None:
"""
- Case 2, the finding-1 regression guard: once execute_sql_statements() has
- opened the connection and asked the (sqlite) engine spec for a cancel
- handle -- and gotten nothing back, since sqlite doesn't override
- get_cancel_query_id() -- a stop request must fail honestly
- (SupersetCancelQueryException, status left alone) instead of fabricating
- a STOPPED status for a query that is, in fact, still running against the
- real database.
-
- get_cancel_query_id() is deliberately left unmocked: the real, permanent
- "this engine has no cancel support" behavior is exactly the case that
- must not be confused with the pre-dispatch race window.
+ Accept Stop after dispatch even when the engine cannot cancel its SQL.
+
+ The real SQLite execute_query still runs after the pause is released, but
+ its result must not be published and its completion cannot replace the
+ accepted Stop or its original timestamp.
"""
from superset import sql_lab
from superset.common.db_query_status import QueryStatus
from superset.constants import QUERY_DISPATCHED_KEY
from superset.daos.query import QueryDAO
- from superset.exceptions import SupersetCancelQueryException
harness = _RaceHarness(mocker)
query_id, client_id = harness.make_query("post-dispatch-no-cancel-support")
@@ -376,15 +357,11 @@ def test_stop_after_dispatch_with_no_cancel_support_raises_honestly(
)
assert "cancel_query" not in query.extra
- # Must raise now: the engine has already been asked for a cancel
- # handle and had none to give, which is a real "can't cancel"
- # failure, not the early race window.
- with pytest.raises(SupersetCancelQueryException):
- QueryDAO.stop_query(client_id)
-
- # Nothing was fabricated: still RUNNING.
- unstopped_query = harness.fresh_query(query_id)
- assert unstopped_query.status == QueryStatus.RUNNING
+ stopped_query = QueryDAO.stop_query(client_id)
+ assert stopped_query.status == QueryStatus.STOPPED
+ stop_end_time = stopped_query.end_time
+ assert stop_end_time is not None
+ assert harness.fresh_query(query_id).status == QueryStatus.STOPPED
finally:
release_execution.set()
@@ -394,15 +371,15 @@ def test_stop_after_dispatch_with_no_cancel_support_raises_honestly(
# The statement genuinely ran (real execute_query was called through).
execute_query_spy.assert_called_once()
- # Since the stop attempt failed, execution proceeds to real completion --
- # this also confirms the success-path finalizer's new STOPPED-guard
- # doesn't misfire and swallow a legitimate SUCCESS.
+ # Engine execution can finish, but an accepted Stop discards its output.
final_query = harness.fresh_query(query_id)
- assert final_query.status == QueryStatus.SUCCESS, (
- f"expected the query to run to real completion since the stop "
- f"attempt correctly failed, got status={final_query.status!r}"
- )
- assert execution_result["payload"]["status"] == QueryStatus.SUCCESS
+ assert final_query.status == QueryStatus.STOPPED
+ assert final_query.end_time == stop_end_time
+ assert final_query.results_key is None
+ assert execution_result["payload"] == {
+ "query_id": query_id,
+ "status": QueryStatus.STOPPED,
+ }
def test_stop_with_existing_cancel_id_unaffected_by_dispatched_key(
diff --git a/tests/unit_tests/daos/test_query_dao.py b/tests/unit_tests/daos/test_query_dao.py
new file mode 100644
index 00000000000..f89d8af7205
--- /dev/null
+++ b/tests/unit_tests/daos/test_query_dao.py
@@ -0,0 +1,608 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for SQL Lab stop requests racing with worker finalization."""
+
+import inspect
+from pathlib import Path
+
+import pytest
+from flask import current_app
+from pytest_mock import MockerFixture
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session
+
+from superset.common.db_query_status import QueryStatus
+from superset.constants import QUERY_EARLY_CANCEL_KEY
+from superset.daos.query import QueryDAO
+from superset.models.core import Database
+from superset.models.sql_lab import Query
+from superset.sql.execution.utils import set_query_stopped
+
+
+@pytest.mark.parametrize(
+ "terminal_status", [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+)
+@pytest.mark.parametrize("early_cancel", [False, True])
+def test_stop_query_preserves_concurrently_committed_terminal_result(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ terminal_status: QueryStatus,
+ early_cancel: bool,
+) -> None:
+ """A worker completes after the stop read but before its guarded write."""
+ engine = create_engine(f"sqlite:///{tmp_path / 'stop.db'}")
+ Query.__table__.create(engine)
+ with Session(engine) as stop_session:
+ query = Query(
+ id=123,
+ client_id="stop-race",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING,
+ extra_json='{"initial": true}',
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ mocker.patch("superset.daos.query.db.session", stop_session)
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+
+ def stop_after_completion(loaded_query: Query, end_time: float) -> bool:
+ """Publish a terminal result after the stop read, before its commit."""
+ assert loaded_query.status == QueryStatus.RUNNING
+ if early_cancel:
+ loaded_query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ with Session(engine) as worker_session:
+ completed = worker_session.get(Query, query.id)
+ assert completed is not None
+ completed.status = terminal_status
+ completed.results_key = "published-results"
+ completed.end_time = 1234.5
+ completed.extra_json = '{"worker": true}'
+ worker_session.commit()
+ return set_query_stopped(loaded_query, end_time)
+
+ mocker.patch(
+ "superset.daos.query.set_query_stopped", side_effect=stop_after_completion
+ )
+ cancel_query = mocker.patch("superset.daos.query.sql_lab.cancel_query")
+ QueryDAO.stop_query("stop-race")
+ cancel_query.assert_not_called()
+
+ with Session(engine) as verify_session:
+ completed = verify_session.get(Query, 123)
+ assert completed is not None
+ assert completed.status == terminal_status
+ assert completed.results_key == "published-results"
+ assert completed.end_time == 1234.5
+ assert completed.extra == {"worker": True}
+ engine.dispose()
+
+
+@pytest.mark.parametrize("completion_timing", ["before_stop", "during_preparation"])
+@pytest.mark.parametrize("engine_name", ["trino", "impala"])
+@pytest.mark.parametrize("caller", ["dao", "async_handle"])
+@pytest.mark.parametrize(
+ "terminal_status", [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+)
+def test_engine_early_cancel_preserves_concurrently_committed_terminal_metadata(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ engine_name: str,
+ caller: str,
+ terminal_status: QueryStatus,
+ completion_timing: str,
+) -> None:
+ """Real engine preparation must not flush stale extra over a worker's result."""
+ from superset.db_engine_specs.impala import ImpalaEngineSpec
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.sql.execution.executor import SQLExecutor
+
+ engine_spec = TrinoEngineSpec if engine_name == "trino" else ImpalaEngineSpec
+ prepare_cancel = engine_spec.prepare_cancel_query
+ mocker.patch.object(Database, "db_engine_spec", engine_spec)
+ engine = create_engine(f"sqlite:///{tmp_path / 'engine-stop.db'}")
+ Database.__table__.create(engine)
+ Database.metadata.tables["ssh_tunnels"].create(engine)
+ Query.__table__.create(engine)
+ with Session(engine) as stop_session:
+ # Core inserts avoid unrelated permission-creation ORM event handlers.
+ stop_session.execute(
+ Database.__table__.insert().values(
+ id=1, database_name="cancel-race", sqlalchemy_uri="sqlite://"
+ )
+ )
+ query = Query(
+ id=123,
+ client_id="engine-stop-race",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING,
+ extra_json='{"initial": true}',
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ mocker.patch("superset.daos.query.db.session", stop_session)
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+
+ def complete() -> None:
+ """Publish an outcome whose target effects finished independently."""
+ with Session(engine) as worker_session:
+ completed = worker_session.get(Query, 123)
+ assert completed is not None
+ completed.status = terminal_status
+ completed.results_key = "published-results"
+ completed.end_time = 1234.5
+ completed.extra_json = '{"columns": [{"name": "worker"}], "errors": []}'
+ worker_session.commit()
+
+ def stop_after_worker_commit(loaded_query: Query, end_time: float) -> bool:
+ """Complete after the stop read, before its initial guarded write."""
+ assert loaded_query.status == QueryStatus.RUNNING
+ complete()
+ return set_query_stopped(loaded_query, end_time)
+
+ def prepare_after_stop(loaded_query: Query) -> None:
+ """A durable result wins while early cancellation is being prepared."""
+ assert loaded_query.status == QueryStatus.STOPPED
+ complete()
+ prepare_cancel(loaded_query)
+
+ prepare = mocker.patch.object(
+ engine_spec, "prepare_cancel_query", side_effect=prepare_after_stop
+ )
+ if completion_timing == "before_stop":
+ mocker.patch(
+ "superset.daos.query.set_query_stopped",
+ side_effect=stop_after_worker_commit,
+ )
+ mocker.patch(
+ "superset.sql.execution.utils.set_query_stopped",
+ side_effect=stop_after_worker_commit,
+ )
+ if caller == "dao":
+ QueryDAO.stop_query("engine-stop-race")
+ else:
+ assert SQLExecutor._cancel_async_query(123, query.database) is False
+ assert prepare.call_count == int(completion_timing == "during_preparation")
+
+ with Session(engine) as verify_session:
+ completed = verify_session.get(Query, 123)
+ assert completed is not None
+ assert completed.status == terminal_status
+ assert completed.results_key == "published-results"
+ assert completed.end_time == 1234.5
+ assert completed.extra == {"columns": [{"name": "worker"}], "errors": []}
+ engine.dispose()
+
+
+@pytest.mark.parametrize("engine_name", ["trino", "impala"])
+@pytest.mark.parametrize("caller", ["dao", "async_handle"])
+def test_engine_early_cancel_commits_with_guarded_stop(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ engine_name: str,
+ caller: str,
+) -> None:
+ """Both cancellation callers publish early intent together with STOPPED."""
+ from superset.db_engine_specs.impala import ImpalaEngineSpec
+ from superset.db_engine_specs.trino import TrinoEngineSpec
+ from superset.sql.execution.executor import SQLExecutor
+
+ engine_spec = TrinoEngineSpec if engine_name == "trino" else ImpalaEngineSpec
+ prepare_cancel = engine_spec.prepare_cancel_query
+ mocker.patch.object(Database, "db_engine_spec", engine_spec)
+ engine = create_engine(f"sqlite:///{tmp_path / 'early-stop.db'}")
+ Database.__table__.create(engine)
+ Database.metadata.tables["ssh_tunnels"].create(engine)
+ Query.__table__.create(engine)
+ with Session(engine) as stop_session:
+ # Core inserts avoid unrelated permission-creation ORM event handlers.
+ stop_session.execute(
+ Database.__table__.insert().values(
+ id=1, database_name="early-cancel", sqlalchemy_uri="sqlite://"
+ )
+ )
+ query = Query(
+ id=123,
+ client_id="early-engine-stop",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING,
+ extra_json='{"initial": true}',
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ mocker.patch("superset.daos.query.db.session", stop_session)
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+
+ def prepare_without_commit(loaded_query: Query) -> None:
+ """Preparation stages early intent after the stop becomes visible."""
+ prepare_cancel(loaded_query)
+ with Session(engine) as observer:
+ stored = observer.get(Query, 123)
+ assert stored is not None
+ assert stored.status == QueryStatus.STOPPED
+ assert stored.extra == {"initial": True}
+
+ mocker.patch.object(
+ engine_spec, "prepare_cancel_query", side_effect=prepare_without_commit
+ )
+ if caller == "dao":
+ QueryDAO.stop_query("early-engine-stop")
+ else:
+ assert SQLExecutor._cancel_async_query(123, query.database) is True
+
+ with Session(engine) as verify_session:
+ stopped = verify_session.get(Query, 123)
+ assert stopped is not None
+ assert stopped.status == QueryStatus.STOPPED
+ assert stopped.extra == {"initial": True, QUERY_EARLY_CANCEL_KEY: True}
+ assert stopped.end_time is not None
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ "terminal_status", [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+)
+def test_stopped_payload_preserves_concurrently_committed_terminal_result(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ terminal_status: QueryStatus,
+) -> None:
+ """A late worker stop payload must not overwrite a published terminal row."""
+ from superset.sql.execution.celery_task import _stopped_query_payload
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'payload.db'}")
+ Query.__table__.create(engine)
+ with Session(engine) as stop_session:
+ query = Query(
+ id=123,
+ client_id="payload-race",
+ database_id=1,
+ status=QueryStatus.STOPPED,
+ end_time=100,
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ assert query.status == QueryStatus.STOPPED
+ with Session(engine) as worker_session:
+ completed = worker_session.get(Query, 123)
+ assert completed is not None
+ completed.status = terminal_status
+ completed.results_key = "published-results"
+ completed.end_time = 1234.5
+ worker_session.commit()
+
+ mocker.patch("superset.sql.execution.celery_task.db.session", stop_session)
+ payload = _stopped_query_payload(query)
+ assert payload["status"] == terminal_status
+
+ with Session(engine) as verify_session:
+ completed = verify_session.get(Query, 123)
+ assert completed is not None
+ assert completed.status == terminal_status
+ assert completed.results_key == "published-results"
+ assert completed.end_time == 1234.5
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ "status",
+ [
+ QueryStatus.PENDING,
+ QueryStatus.RUNNING,
+ QueryStatus.SCHEDULED,
+ QueryStatus.FETCHING,
+ QueryStatus.STOPPED,
+ ],
+)
+def test_stop_query_commits_stop_with_early_cancel(
+ mocker: MockerFixture, app_context: None, session: Session, status: QueryStatus
+) -> None:
+ """Only the first Stop stages the early-cancel flag; STOPPED suffices after it."""
+ Query.__table__.create(session.get_bind())
+ query = Query(client_id="early-stop", user_id=7, database_id=1, status=status)
+ session.add(query)
+ session.commit()
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ mocker.patch("superset.daos.query.now_as_float", return_value=1234.5)
+
+ def cancel(loaded_query: Query) -> bool:
+ """Stage stop intent without exposing it separately from STOPPED."""
+ loaded_query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ stored_extra = (
+ session.query(Query.extra_json).filter(Query.id == query.id).scalar()
+ )
+ assert stored_extra == "{}"
+ return True
+
+ mocker.patch("superset.daos.query.sql_lab.cancel_query", side_effect=cancel)
+ QueryDAO.stop_query("early-stop")
+
+ session.expire_all()
+ assert query.status == QueryStatus.STOPPED
+ assert bool(query.extra.get(QUERY_EARLY_CANCEL_KEY)) is (
+ status != QueryStatus.STOPPED
+ )
+ assert query.end_time == 1234.5
+
+
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("finish_before_stop", [False, True])
+@pytest.mark.parametrize(
+ "terminal_status", [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+)
+def test_stop_reports_concurrently_finished_query(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ caller: str,
+ finish_before_stop: bool,
+ terminal_status: QueryStatus,
+) -> None:
+ """A finished query is not reported as cancelled, including a lost stop race."""
+ from superset.queries.api import QueryRestApi
+ from superset.sql.execution.executor import SQLExecutor
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'stop-outcome.db'}")
+ Query.__table__.create(engine)
+ with Session(engine) as stop_session:
+ query = Query(
+ id=123,
+ client_id="finished-stop",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING if finish_before_stop else terminal_status,
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ mocker.patch("superset.daos.query.db.session", stop_session)
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ payload = {"id": "finished-stop", "state": terminal_status.lower()}
+ mocker.patch.object(Query, "to_dict", return_value=payload)
+
+ def stop_after_completion(loaded_query: Query, end_time: float) -> bool:
+ """Commit the terminal outcome before the guarded stop write."""
+ assert finish_before_stop
+ with Session(engine) as worker_session:
+ completed = worker_session.get(Query, 123)
+ assert completed is not None
+ completed.status = terminal_status
+ worker_session.commit()
+ return set_query_stopped(loaded_query, end_time)
+
+ mocker.patch(
+ "superset.daos.query.set_query_stopped", side_effect=stop_after_completion
+ )
+ mocker.patch(
+ "superset.sql.execution.utils.set_query_stopped",
+ side_effect=stop_after_completion,
+ )
+ cancel_query = mocker.patch("superset.daos.query.sql_lab.cancel_query")
+ cancel_async = mocker.patch.object(SQLExecutor, "_cancel_query")
+ if caller == "api":
+ with current_app.test_request_context(json={"client_id": "finished-stop"}):
+ response = inspect.unwrap(QueryRestApi.stop_query)(QueryRestApi())
+ assert response.status_code == 422
+ assert response.get_json() == {
+ "message": "Could not cancel query",
+ "query": payload,
+ }
+ else:
+ assert SQLExecutor._cancel_async_query(123, mocker.Mock()) is False
+ cancel_query.assert_not_called()
+ cancel_async.assert_not_called()
+ assert stop_session.get(Query, 123).status == terminal_status
+ engine.dispose()
+
+
+@pytest.mark.parametrize("caller", ["dao", "async_handle"])
+@pytest.mark.parametrize("early_cancel", [False, True])
+def test_winning_stop_preserves_concurrently_committed_extra(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ caller: str,
+ early_cancel: bool,
+) -> None:
+ """Only stop-owned metadata changes when a worker updates a running query."""
+ from superset.constants import QUERY_CANCEL_KEY
+ from superset.sql.execution.executor import SQLExecutor
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'stop-extra.db'}")
+ Query.__table__.create(engine)
+ worker_extra = {QUERY_CANCEL_KEY: "worker-handle", "progress": "worker progress"}
+ with Session(engine) as stop_session:
+ query = Query(
+ id=123,
+ client_id="extra-stop",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING,
+ extra_json='{"initial": true}',
+ )
+ stop_session.add(query)
+ stop_session.commit()
+ mocker.patch("superset.daos.query.db.session", stop_session)
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+
+ def cancel(*args: object) -> bool:
+ """Stage early intent while a second session publishes worker metadata."""
+ assert query.extra == {"initial": True}
+ if early_cancel:
+ query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ with Session(engine) as worker_session:
+ running = worker_session.get(Query, 123)
+ assert running is not None
+ for key, value in worker_extra.items():
+ running.set_extra_json_key(key, value)
+ worker_session.commit()
+ return True
+
+ mocker.patch("superset.daos.query.sql_lab.cancel_query", side_effect=cancel)
+ mocker.patch.object(SQLExecutor, "_cancel_query", side_effect=cancel)
+ if caller == "dao":
+ QueryDAO.stop_query("extra-stop")
+ else:
+ assert SQLExecutor._cancel_async_query(123, mocker.Mock()) is True
+
+ with Session(engine) as verify_session:
+ stopped = verify_session.get(Query, 123)
+ assert stopped is not None
+ assert stopped.status == QueryStatus.STOPPED
+ assert stopped.end_time is not None
+ assert stopped.extra == {
+ "initial": True,
+ **worker_extra,
+ **({QUERY_EARLY_CANCEL_KEY: True} if early_cancel else {}),
+ }
+ engine.dispose()
+
+
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("cancel_outcome", ["unsupported", "unreachable"])
+@pytest.mark.parametrize("status", [QueryStatus.PENDING, QueryStatus.RUNNING])
+def test_accepted_stop_reports_success_when_engine_cancel_fails(
+ mocker: MockerFixture,
+ app_context: None,
+ session: Session,
+ caller: str,
+ cancel_outcome: str,
+ status: QueryStatus,
+ caplog: pytest.LogCaptureFixture,
+) -> None:
+ """A committed stop is accepted even if engine cancellation is unavailable."""
+ from superset.queries.api import QueryRestApi
+ from superset.sql.execution.executor import SQLExecutor
+
+ Query.__table__.create(session.get_bind())
+ query = Query(
+ id=123,
+ client_id="accept-stop",
+ user_id=7,
+ database_id=1,
+ status=status,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ cancel = mocker.Mock(return_value=False)
+ if cancel_outcome == "unreachable":
+ cancel.side_effect = RuntimeError("Engine unavailable")
+ mocker.patch("superset.daos.query.sql_lab.cancel_query", cancel)
+ mocker.patch.object(SQLExecutor, "_cancel_query", cancel)
+
+ if caller == "api":
+ with current_app.test_request_context(json={"client_id": "accept-stop"}):
+ response = inspect.unwrap(QueryRestApi.stop_query)(QueryRestApi())
+ assert response.status_code == 200
+ assert response.get_json() == {"result": "OK"}
+ else:
+ handle = SQLExecutor(mocker.Mock())._create_async_handle(123)
+ assert handle.cancel() is True
+ cancel.assert_called_once()
+ warning = next(
+ record
+ for record in caplog.records
+ if "engine cancellation not confirmed" in record.message
+ )
+ assert bool(warning.exc_info) is (cancel_outcome == "unreachable")
+ session.expire_all()
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time is not None
+
+
+@pytest.mark.parametrize("caller", ["api", "async_handle"])
+@pytest.mark.parametrize("failure", ["none", "final_stop", "final_commit"])
+def test_stop_acceptance_survives_post_commit_failure(
+ mocker: MockerFixture,
+ app_context: None,
+ session: Session,
+ caller: str,
+ failure: str,
+) -> None:
+ """A durable stop remains accepted, and repeating it preserves its timestamp."""
+ from superset.queries.api import QueryRestApi
+ from superset.sql.execution.executor import SQLExecutor
+
+ Query.__table__.create(session.get_bind())
+ query = Query(
+ id=123,
+ client_id="durable-stop",
+ user_id=7,
+ database_id=1,
+ status=QueryStatus.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ cancel_api = mocker.patch(
+ "superset.daos.query.sql_lab.cancel_query", return_value=False
+ )
+ cancel_handle = mocker.patch.object(
+ SQLExecutor, "_cancel_query", return_value=False
+ )
+ commit = session.commit
+ stop = set_query_stopped
+ commits = 0
+ stops = 0
+
+ def fail_commit() -> None:
+ """Fail only the second commit, after acceptance is durable."""
+ nonlocal commits
+ commits += 1
+ if failure == "final_commit" and commits == 2:
+ raise RuntimeError("post-stop commit failed")
+ commit()
+
+ def fail_stop(query: Query, end_time: float) -> bool:
+ """Fail the post-cancellation metadata merge."""
+ nonlocal stops
+ stops += 1
+ if failure == "final_stop" and stops == 2:
+ raise RuntimeError("post-stop merge failed")
+ return stop(query, end_time)
+
+ mocker.patch.object(session, "commit", side_effect=fail_commit)
+ mocker.patch("superset.daos.query.set_query_stopped", side_effect=fail_stop)
+ mocker.patch(
+ "superset.sql.execution.utils.set_query_stopped", side_effect=fail_stop
+ )
+
+ def request_stop() -> None:
+ """Exercise the public response rather than only the staged row."""
+ if caller == "api":
+ with current_app.test_request_context(json={"client_id": "durable-stop"}):
+ response = inspect.unwrap(QueryRestApi.stop_query)(QueryRestApi())
+ assert response.status_code == 200
+ else:
+ assert SQLExecutor(mocker.Mock())._create_async_handle(123).cancel() is True
+
+ request_stop()
+ session.expire_all()
+ original_end_time = query.end_time
+ assert query.status == QueryStatus.STOPPED
+ request_stop()
+ session.expire_all()
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time == original_end_time
+ # Stored cancellation handles may identify an unrelated session after reuse.
+ (cancel_api if caller == "api" else cancel_handle).assert_called_once()
diff --git a/tests/unit_tests/db_engine_specs/test_impala_pending.py b/tests/unit_tests/db_engine_specs/test_impala_pending.py
index e656113d1dd..b8cef54a868 100644
--- a/tests/unit_tests/db_engine_specs/test_impala_pending.py
+++ b/tests/unit_tests/db_engine_specs/test_impala_pending.py
@@ -164,7 +164,8 @@ def test_stop_before_cancel_id_is_published_uses_live_cursor(
assert SQLExecutor._cancel_query(query.database, query)
else:
assert cancel_query(query)
- db.session.commit.assert_called_once_with()
+ # The caller commits early intent together with the guarded STOPPED write.
+ db.session.commit.assert_not_called()
cursor: Mock = Mock()
cursor.status.return_value = "PENDING_STATE"
diff --git a/tests/unit_tests/sql/execution/conftest.py b/tests/unit_tests/sql/execution/conftest.py
index 296ae8f8835..4aa29606330 100644
--- a/tests/unit_tests/sql/execution/conftest.py
+++ b/tests/unit_tests/sql/execution/conftest.py
@@ -100,6 +100,7 @@ def mock_database() -> MagicMock:
database.db_engine_spec.get_cancel_query_id = MagicMock(return_value=None)
database.db_engine_spec.patch = MagicMock()
database.db_engine_spec.fetch_data = MagicMock(return_value=[])
+ database.db_engine_spec.is_autocommit = MagicMock(return_value=False)
# Mirrors the real `Database.mutate_sql_based_on_config` default (no-op
# when no `SQL_QUERY_MUTATOR` is configured), so SQL parsed from its
# return value stays valid instead of an un-parseable `MagicMock`.
diff --git a/tests/unit_tests/sql/execution/test_celery_task.py b/tests/unit_tests/sql/execution/test_celery_task.py
index fc12ae2616e..5876b8d3ab4 100644
--- a/tests/unit_tests/sql/execution/test_celery_task.py
+++ b/tests/unit_tests/sql/execution/test_celery_task.py
@@ -17,9 +17,11 @@
"""Tests for celery_task.py - async SQL execution via Celery."""
+from collections.abc import Callable
from decimal import Decimal
+from pathlib import Path
from typing import Any
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, PropertyMock
import msgpack
import pandas as pd
@@ -1143,6 +1145,135 @@ def test_execute_sql_task_stopped(
assert result["status"] == QueryStatusEnum.STOPPED
+@pytest.mark.parametrize("stop_intent", ["early", "status"])
+def test_cancelled_impala_cursor_is_not_finalized_as_success(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_database: MagicMock,
+ stop_intent: str,
+) -> None:
+ """Cleared impyla description after a stop is not a successful DML result."""
+ impyla = pytest.importorskip("impala.hiveserver2")
+
+ from superset.constants import QUERY_CANCEL_KEY, QUERY_EARLY_CANCEL_KEY
+ from superset.db_engine_specs.impala import ImpalaEngineSpec
+ from superset.sql.execution import celery_task
+
+ mock_query.sql = "SELECT 1"
+ mock_query.database = mock_database
+ mock_database.db_engine_spec = ImpalaEngineSpec
+ mock_query.extra = {QUERY_CANCEL_KEY: "0123456789abcdef:fedcba9876543210"}
+ mock_query.set_extra_json_key.side_effect = lambda key, value: (
+ mock_query.extra.update({key: value})
+ )
+ cursor = impyla.HiveServer2Cursor(MagicMock())
+ operation = MagicMock(has_result_set=True)
+
+ def execute_async(sql: str) -> None:
+ """Publish a real impyla operation backed by mocked RPCs."""
+ cursor._last_operation = operation
+ cursor._last_operation_active = True
+
+ def stop_before_poll() -> str:
+ """Stop after dispatch, before handle_cursor checks the intent."""
+ if stop_intent == "early":
+ mock_query.extra[QUERY_EARLY_CANCEL_KEY] = True
+ else:
+ mock_query.status = QueryStatusEnum.STOPPED
+ return "RUNNING_STATE"
+
+ operation.get_status.side_effect = stop_before_poll
+ mocker.patch.object(cursor, "execute_async", side_effect=execute_async)
+ connection = mock_database.get_raw_connection.return_value.__enter__.return_value
+ connection.cursor.return_value = cursor
+ mocker.patch.object(ImpalaEngineSpec, "get_cancel_query_id", return_value=None)
+ mocker.patch(
+ "superset.sql.execution.celery_task._get_query", return_value=mock_query
+ )
+ finalize = mocker.spy(celery_task, "_finalize_successful_query")
+ store = mocker.patch("superset.sql.execution.celery_task._store_results_in_backend")
+ mocker.patch("superset.sql.execution.celery_task.results_backend", MagicMock())
+ post = mocker.patch("superset.db_engine_specs.impala.requests.post")
+
+ payload = celery_task._execute_sql_statements(123, "SELECT 1", None)
+
+ operation.cancel.assert_called_once_with()
+ operation.close.assert_called_once_with()
+ assert cursor.description is None
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ assert mock_query.status == QueryStatusEnum.STOPPED
+ connection.commit.assert_not_called()
+ finalize.assert_not_called()
+ store.assert_not_called()
+ post.assert_not_called()
+
+
+@pytest.mark.parametrize("stop_intent", ["early", "status"])
+@pytest.mark.parametrize("stop_checkpoint", ["execution", "connection_exit"])
+@pytest.mark.parametrize("end_time", [None, 1234.5])
+def test_stop_with_partial_results_skips_success_finalization(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_database: MagicMock,
+ stop_intent: str,
+ stop_checkpoint: str,
+ end_time: float | None,
+) -> None:
+ """
+ A stop seen by the pre-commit check skips success finalization and keeps
+ the recorded end time. A stop during connection cleanup is caught by the
+ guarded publish instead, after the results are prepared.
+ """
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.sql.execution import celery_task
+
+ mock_query.database = mock_database
+ mock_query.end_time = end_time
+ mocker.patch.object(celery_task, "_get_query", return_value=mock_query)
+ mocker.patch.object(celery_task, "now_as_float", return_value=5678.9)
+ finalize = mocker.spy(celery_task, "_finalize_successful_query")
+ store = mocker.patch.object(celery_task, "_store_results_in_backend")
+ mocker.patch.object(celery_task, "results_backend", MagicMock())
+
+ def record_stop(*args: Any, **kwargs: Any) -> None:
+ """Expose the stop intent at the selected execution boundary."""
+ if stop_intent == "early":
+ mock_query.extra[QUERY_EARLY_CANCEL_KEY] = True
+ else:
+ mock_query.status = QueryStatusEnum.STOPPED
+
+ results = [("SELECT * FROM users", None, 1.0, 0)]
+
+ def execute_with_partial_results(
+ **kwargs: Any,
+ ) -> list[tuple[str, None, float, int]]:
+ """Return completed work even if the last statement was stopped."""
+ if stop_checkpoint == "execution":
+ record_stop()
+ return results
+
+ mocker.patch.object(
+ celery_task, "execute_sql_with_cursor", side_effect=execute_with_partial_results
+ )
+ if stop_checkpoint == "connection_exit":
+ mock_database.get_raw_connection.return_value.__exit__.side_effect = record_stop
+
+ payload = celery_task._execute_sql_statements(123, mock_query.sql, None)
+
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ assert mock_query.status == QueryStatusEnum.STOPPED
+ if stop_checkpoint == "execution":
+ assert mock_query.end_time == (end_time if end_time is not None else 5678.9)
+ finalize.assert_not_called()
+ store.assert_not_called()
+ else:
+ finalize.assert_called_once()
+ store.assert_called_once()
+
+
def test_execute_sql_task_with_mutation(
mocker: MockerFixture,
app_context: None,
@@ -1347,7 +1478,7 @@ def test_execute_sql_task_with_failed_status_before_final_commit(
)
# Mock _store_results_in_backend to set status to FAILED without raising
- def mock_store_results(query, payload, database):
+ def mock_store_results(query, payload, database, on_results_key=None):
query.status = QueryStatusEnum.FAILED
mocker.patch("superset.sql.execution.celery_task.results_backend", MagicMock())
@@ -1369,4 +1500,2088 @@ def test_execute_sql_task_with_failed_status_before_final_commit(
# Verify query status remains FAILED and is not changed to SUCCESS
assert result is not None
+ assert "error" not in result
assert mock_query.status == QueryStatusEnum.FAILED
+
+
+@pytest.mark.parametrize(
+ "stop_checkpoint",
+ ["serialization", "backend_write", "backend_failure", "backend_error"],
+)
+def test_stop_during_result_finalization_stays_stopped(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_db_session: MagicMock,
+ mock_query: MagicMock,
+ mock_database: MagicMock,
+ mock_result_set: MagicMock,
+ stop_checkpoint: str,
+) -> None:
+ """A stop during result processing must not become a successful query."""
+ from superset.sql.execution import celery_task
+
+ mock_query.database = mock_database
+ mocker.patch.object(celery_task, "_get_query", return_value=mock_query)
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[(mock_query.sql, mock_result_set, 1.0, 2)],
+ )
+ backend = mocker.patch.object(celery_task, "results_backend", MagicMock())
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"results")
+
+ committed: dict[str, Any] = {}
+
+ def record_stop() -> None:
+ """Another request commits a stop after the post-execution check."""
+ committed.update(status=QueryStatusEnum.STOPPED, end_time=1234.5)
+ mock_query.status = QueryStatusEnum.STOPPED
+
+ def refresh(query: MagicMock, *args: Any, **kwargs: Any) -> None:
+ """A refresh reloads the committed stop over pending success fields."""
+ for key, value in committed.items():
+ setattr(query, key, value)
+
+ mock_db_session.refresh.side_effect = refresh
+
+ def serialize(result_set: MagicMock) -> tuple[list[object], list[object]]:
+ """Stop while preparing the success payload."""
+ if stop_checkpoint == "serialization":
+ record_stop()
+ return [], []
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Stop while the results backend write is in flight."""
+ if stop_checkpoint != "serialization":
+ record_stop()
+ if stop_checkpoint == "backend_error":
+ raise RuntimeError("Backend unavailable")
+ return stop_checkpoint != "backend_failure"
+
+ mocker.patch.object(celery_task, "_serialize_result_set", side_effect=serialize)
+ backend.set.side_effect = store
+
+ mocker.patch.object(celery_task, "security_manager", MagicMock())
+ payload = celery_task.execute_sql_task(123, mock_query.sql)
+
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ assert mock_query.status == QueryStatusEnum.STOPPED
+ assert mock_query.end_time == 1234.5
+ assert mock_query.results_key is None
+ # A stop during serialization is caught after the backend write.
+ backend.set.assert_called_once()
+
+
+@pytest.mark.parametrize("delete_fails", [False, True], ids=["deleted", "delete_error"])
+@pytest.mark.parametrize(
+ "checkpoint",
+ ["none", "backend_write", "backend_failure", "publish_check", "publish_commit"],
+)
+def test_stop_after_results_write_deletes_stored_results(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ caplog: pytest.LogCaptureFixture,
+ checkpoint: str,
+ delete_fails: bool,
+) -> None:
+ """
+ Test that results stored for a query a stop then wins are deleted from the
+ results backend, and that a failed delete still ends STOPPED.
+ """
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'stored.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ def deliver_stop() -> None:
+ """Commit the stop in an independent transaction."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob, with the stop landing while the write is in flight."""
+ if checkpoint in {"backend_write", "backend_failure"}:
+ deliver_stop()
+ if checkpoint == "backend_failure":
+ return False
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ # A failing delete behaves like an unavailable backend.
+ backend.delete.side_effect = (
+ RuntimeError("Backend unavailable") if delete_fails else delete
+ )
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+
+ original_make_check = celery_task._make_check_stopped_fn
+ finalization_checks = 0
+
+ def make_check(query: Query, preserve_pending: bool = False) -> Callable[[], bool]:
+ """Deliver Stop after the post-write check or after the publish check."""
+ check = original_make_check(query, preserve_pending)
+
+ def check_then_stop() -> bool:
+ nonlocal finalization_checks
+ stopped = check()
+ if preserve_pending:
+ finalization_checks += 1
+ if (checkpoint, finalization_checks) in {
+ ("publish_check", 1),
+ ("publish_commit", 2),
+ }:
+ deliver_stop()
+ return stopped
+
+ return check_then_stop
+
+ mocker.patch.object(celery_task, "_make_check_stopped_fn", side_effect=make_check)
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="stored-results",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ session.refresh(query)
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+
+ with caplog.at_level("WARNING", logger=celery_task.logger.name):
+ payload = celery_task._execute_sql_statements(123, query.sql, None)
+
+ session.refresh(query)
+ _assert_stored_results_outcome(
+ payload, query, backend, stored, caplog.text, checkpoint, delete_fails
+ )
+ engine.dispose()
+
+
+def _assert_stored_results_outcome(
+ payload: dict[str, Any] | None,
+ query: Any,
+ backend: MagicMock,
+ stored: dict[str, bytes],
+ log_text: str,
+ checkpoint: str,
+ delete_fails: bool,
+) -> None:
+ """Check the query, payload and results backend after finalization."""
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ if checkpoint == "none":
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.SUCCESS
+ assert query.status == QueryStatusEnum.SUCCESS
+ assert query.results_key == key
+ assert key in stored
+ backend.delete.assert_not_called()
+ return
+
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.end_time == Decimal("1234.5")
+ assert query.results_key is None
+ if checkpoint == "backend_failure":
+ # Nothing was stored, so there is nothing to delete.
+ backend.delete.assert_not_called()
+ return
+
+ backend.delete.assert_called_once_with(key)
+ assert (key in stored) is delete_fails
+ assert ("Failed to delete unpublished results" in log_text) is delete_fails
+
+
+@pytest.mark.parametrize("delete_fails", [False, True], ids=["deleted", "delete_error"])
+@pytest.mark.parametrize(
+ "failure,stop",
+ [
+ (failure, stop)
+ for failure in ["write_timing", "status_read", "publish_update"]
+ for stop in ["none", "with_failure", "error_commit"]
+ ]
+ # The failed flush follows the guarded UPDATE, whose SQLite write lock
+ # blocks a stop committed from another connection until the rollback.
+ + [("publish_flush", "none"), ("publish_flush", "error_commit")],
+)
+def test_error_after_results_write_deletes_stored_results_when_stopped(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ caplog: pytest.LogCaptureFixture,
+ failure: str,
+ stop: str,
+ delete_fails: bool,
+) -> None:
+ """
+ Test that results stored before a task error are deleted when the error
+ handler finds the query stopped, and that a failed delete still ends STOPPED.
+ """
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'error-stored.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ def deliver_stop() -> None:
+ """Commit the stop in an independent transaction."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = (
+ RuntimeError("Backend unavailable") if delete_fails else delete
+ )
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ failed = False
+
+ def fail_after_write() -> None:
+ """Raise once the results are stored, with the stop landing first."""
+ nonlocal failed
+ assert stored
+ failed = True
+ if stop == "with_failure":
+ deliver_stop()
+ raise RuntimeError("Metrics or metadata database unavailable")
+
+ _fail_after_results_write(
+ mocker, engine, failure, lambda: bool(stored) and not failed, fail_after_write
+ )
+
+ original_make_check = celery_task._make_check_stopped_fn
+
+ def make_check(query: Query, preserve_pending: bool = False) -> Callable[[], bool]:
+ """Deliver the stop after the error handler's own stop check."""
+ check = original_make_check(query, preserve_pending)
+
+ def check_then_stop() -> bool:
+ stopped = check()
+ if stop == "error_commit" and failed and preserve_pending:
+ deliver_stop()
+ return stopped
+
+ return check_then_stop
+
+ mocker.patch.object(celery_task, "_make_check_stopped_fn", side_effect=make_check)
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="error-stored",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ with caplog.at_level("WARNING", logger=celery_task.logger.name):
+ payload = celery_task.execute_sql_task(123, query.sql)
+
+ session.refresh(query)
+ assert failed
+ _assert_error_results_outcome(
+ payload, query, backend, stored, caplog.text, stop, delete_fails
+ )
+ engine.dispose()
+
+
+def _fail_after_results_write(
+ mocker: MockerFixture,
+ engine: Any,
+ failure: str,
+ armed: Callable[[], bool],
+ fail: Callable[[], None],
+) -> None:
+ """
+ Raise from the write timing metric, the status read right after the write,
+ the guarded publish UPDATE, or the flush committing the published metadata.
+ """
+ from sqlalchemy import event
+
+ if failure == "write_timing":
+ stats_logger = MagicMock()
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the results backend write timing."""
+ if key == "sqllab.query.results_backend_write" and armed():
+ fail()
+
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(current_app.config, {"STATS_LOGGER": stats_logger})
+ return
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def fail_publish( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Fail the selected metadata statement, after the results are stored."""
+ verb = statement.lstrip().split(" ", 1)[0]
+ guarded = "status !=" in statement
+ failing = {
+ "status_read": verb == "SELECT",
+ "publish_update": verb == "UPDATE" and guarded,
+ "publish_flush": verb == "UPDATE" and not guarded,
+ }[failure]
+ if failing and armed():
+ fail()
+
+
+def _assert_error_results_outcome(
+ payload: dict[str, Any] | None,
+ query: Any,
+ backend: MagicMock,
+ stored: dict[str, bytes],
+ log_text: str,
+ stop: str,
+ delete_fails: bool,
+) -> None:
+ """Check the query, payload and results backend after the task error."""
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ if stop == "none":
+ # Not a stop: the failure is reported and the row keeps its key.
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.FAILED
+ assert query.status == QueryStatusEnum.FAILED
+ assert query.results_key == key
+ assert key in stored
+ backend.delete.assert_not_called()
+ return
+
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.end_time == Decimal("1234.5")
+ assert query.results_key is None
+ backend.delete.assert_called_once_with(key)
+ assert (key in stored) is delete_fails
+ assert ("Failed to delete unpublished results" in log_text) is delete_fails
+
+
+@pytest.mark.parametrize(
+ ("delete_result", "message"),
+ [
+ (True, None),
+ (False, "Unpublished results were not deleted"),
+ (RuntimeError("Backend unavailable"), "Failed to delete unpublished results"),
+ ],
+ ids=["deleted", "missing", "delete_error"],
+)
+def test_discard_stored_results_logs_undeleted_results(
+ mocker: MockerFixture,
+ app_context: None,
+ caplog: pytest.LogCaptureFixture,
+ delete_result: bool | Exception,
+ message: str | None,
+) -> None:
+ """Test that results the backend did not delete are logged, not raised."""
+ from superset.sql.execution import celery_task
+
+ backend = MagicMock()
+ if isinstance(delete_result, Exception):
+ backend.delete.side_effect = delete_result
+ else:
+ backend.delete.return_value = delete_result
+ mocker.patch.object(celery_task, "results_backend", backend)
+
+ with caplog.at_level("WARNING", logger="superset.sql.execution.utils"):
+ celery_task._discard_stored_results(123, "results-key", backend)
+
+ backend.delete.assert_called_once_with("results-key")
+ if message is None:
+ assert not caplog.records
+ else:
+ assert message in caplog.text
+
+
+@pytest.mark.parametrize(
+ ("key", "has_backend"),
+ [(None, True), ("results-key", False)],
+ ids=["no_key", "no_backend"],
+)
+def test_discard_unreferenced_results_without_stored_results(
+ mocker: MockerFixture,
+ app_context: None,
+ key: str | None,
+ has_backend: bool,
+) -> None:
+ """Test that nothing is read or deleted when no results were stored."""
+ from superset.sql.execution import celery_task
+
+ backend = MagicMock() if has_backend else None
+ mocker.patch.object(celery_task, "results_backend", backend)
+ session = mocker.patch.object(celery_task.db, "session")
+
+ celery_task._discard_unreferenced_results(123, key, backend)
+
+ session.rollback.assert_not_called()
+ session.query.assert_not_called()
+ if backend is not None:
+ backend.delete.assert_not_called()
+
+
+@pytest.mark.parametrize("failure", ["payload", "commit"])
+def test_error_path_deletes_results_only_without_committed_failure(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ failure: str,
+) -> None:
+ """
+ Test that the error path deletes results stored before a failure only when
+ no committed row references them: a failure building the error payload or
+ committing the failed row leaves the row uncommitted, so its results go.
+ """
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'failed-commit.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = delete
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+
+ stats_logger = MagicMock()
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the write timing, once the results are stored."""
+ if key == "sqllab.query.results_backend_write":
+ raise RuntimeError("Metrics unavailable")
+
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(
+ current_app.config, {"STATS_LOGGER": stats_logger, "SHOW_STACKTRACE": True}
+ )
+ if failure == "payload":
+ mocker.patch.object(
+ celery_task.traceback,
+ "format_exc",
+ side_effect=RuntimeError("Error payload failed"),
+ )
+
+ failed_flushed = False
+
+ @event.listens_for(engine, "commit")
+ def fail_failed_commit(conn: Any) -> None:
+ """Fail committing the failed row before the database commits it."""
+ if failure == "commit" and failed_flushed:
+ raise RuntimeError("Metadata commit failed")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="failed-commit",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ @event.listens_for(session, "after_flush")
+ def flushed_failure(flushing_session: Any, context: Any) -> None:
+ nonlocal failed_flushed
+ failed_flushed = query.status == QueryStatusEnum.FAILED
+
+ match = "Error payload failed" if failure == "payload" else "commit failed"
+ with pytest.raises(RuntimeError, match=match):
+ celery_task.execute_sql_task(123, query.sql)
+ session.rollback()
+
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatusEnum.RUNNING
+ assert persisted.results_key is None
+ backend.delete.assert_called_once_with(key)
+ assert key not in stored
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ ("reread", "delete_fails"),
+ [("available", False), ("available", True), ("unavailable", False)],
+ ids=["reread-deleted", "reread-delete_error", "no_reread"],
+)
+def test_unrecorded_error_after_results_write_deletes_only_unreferenced_results(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ caplog: pytest.LogCaptureFixture,
+ reread: str,
+ delete_fails: bool,
+) -> None:
+ """
+ Test that results stored before a failure the task cannot record are
+ deleted once a fresh read shows the row does not reference them, and kept
+ until they expire when the metadata database cannot be read at all.
+ """
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'unrecorded.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = (
+ RuntimeError("Backend unavailable") if delete_fails else delete
+ )
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def metadata_unavailable( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Fail every metadata statement once the results are stored."""
+ cleanup_read = statement.startswith('SELECT "query".results_key ')
+ if stored and not (reread == "available" and cleanup_read):
+ raise RuntimeError("Metadata database unavailable")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="unrecorded",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ with caplog.at_level("WARNING", logger=celery_task.logger.name):
+ with pytest.raises(RuntimeError, match="Metadata database unavailable"):
+ celery_task.execute_sql_task(123, query.sql)
+
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ if reread == "available":
+ backend.delete.assert_called_once_with(key)
+ assert (key in stored) is delete_fails
+ assert (
+ "Failed to delete unpublished results" in caplog.text
+ ) is delete_fails
+ else:
+ backend.delete.assert_not_called()
+ assert key in stored
+ assert "Kept results that may be referenced" in caplog.text
+ engine.dispose()
+
+
+@pytest.mark.parametrize("applied", ["failed_row", "published_row"])
+def test_commit_that_applies_and_raises_keeps_referenced_results(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ applied: str,
+) -> None:
+ """
+ Test that a metadata commit that reaches the database and then raises (its
+ acknowledgement is lost) does not delete the results the committed row
+ references: the failed row of the error handler, or the published row
+ when the metadata database then stays unavailable.
+ """
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'applied-commit.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ stats_logger = MagicMock()
+ if applied == "failed_row":
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the write timing, once the results are stored."""
+ if key == "sqllab.query.results_backend_write":
+ raise RuntimeError("Metrics unavailable")
+
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(current_app.config, {"STATS_LOGGER": stats_logger})
+
+ target = QueryStatusEnum.FAILED if applied == "failed_row" else "success"
+ flushed = False
+ applied_commit = False
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def metadata_unavailable( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Keep the metadata database down once the published row commits."""
+ if applied == "published_row" and applied_commit:
+ raise RuntimeError("Metadata database unavailable")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="applied-commit",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ @event.listens_for(session, "after_flush")
+ def flushed_target(flushing_session: Any, context: Any) -> None:
+ nonlocal flushed
+ flushed = query.status == target and query.results_key is not None
+
+ do_commit = engine.dialect.do_commit
+
+ def lose_acknowledgement(dbapi_connection: Any) -> None:
+ """Raise after the database committed the row with the key."""
+ nonlocal applied_commit
+ do_commit(dbapi_connection)
+ if flushed and not applied_commit:
+ applied_commit = True
+ raise RuntimeError("Commit acknowledgement lost")
+
+ mocker.patch.object(
+ engine.dialect, "do_commit", side_effect=lose_acknowledgement
+ )
+
+ with pytest.raises(RuntimeError):
+ celery_task.execute_sql_task(123, query.sql)
+
+ assert applied_commit
+ event.remove(engine, "before_cursor_execute", metadata_unavailable)
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == target
+ assert persisted.results_key == key
+ backend.delete.assert_not_called()
+ assert key in stored
+ engine.dispose()
+
+
+@pytest.mark.parametrize("stop_at", ["before_stop_check", "before_failed_commit"])
+def test_stop_after_an_applied_publish_keeps_referenced_results(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ stop_at: str,
+) -> None:
+ """
+ Test that when the publish commit reaches the database and then raises, and
+ a stop lands while the error is handled, the stop wins without deleting the
+ results the committed row references, whichever stop check sees it.
+ """
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'applied-publish.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ mocker.patch.dict(current_app.config, {"STATS_LOGGER": MagicMock()})
+
+ def stop() -> None:
+ """Commit a stop from another session, as the stop request does."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stop_session.commit()
+
+ def extract_errors(*args: Any, **kwargs: Any) -> list[Any]:
+ """Stop after the handler's first check, before its guarded commit."""
+ stop()
+ return []
+
+ if stop_at == "before_failed_commit":
+ mock_database.db_engine_spec.extract_errors.side_effect = extract_errors
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="applied-publish",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ published = False
+
+ @event.listens_for(session, "after_flush")
+ def published_flush(flushing_session: Any, context: Any) -> None:
+ nonlocal published
+ published = query.results_key is not None
+
+ do_commit = engine.dialect.do_commit
+ applied_commit = False
+
+ def lose_acknowledgement(dbapi_connection: Any) -> None:
+ """Raise after the database committed the published row."""
+ nonlocal applied_commit
+ do_commit(dbapi_connection)
+ if published and not applied_commit:
+ applied_commit = True
+ if stop_at == "before_stop_check":
+ stop()
+ raise RuntimeError("Commit acknowledgement lost")
+
+ mocker.patch.object(
+ engine.dialect, "do_commit", side_effect=lose_acknowledgement
+ )
+
+ payload = celery_task.execute_sql_task(123, query.sql)
+
+ assert applied_commit
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatusEnum.STOPPED
+ assert persisted.results_key == key
+ backend.delete.assert_not_called()
+ assert key in stored
+ engine.dispose()
+
+
+@pytest.mark.parametrize("stopped", [False, True])
+def test_finalization_stop_check_preserves_pending_metadata(
+ mocker: MockerFixture, app_context: None, stopped: bool
+) -> None:
+ """A status-only refresh keeps success fields but discards them on a stop."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine("sqlite://")
+ Query.__table__.create(engine)
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="stop-test",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ session.refresh(query)
+ mocker.patch.object(celery_task.db, "session", session)
+ if stopped:
+ with engine.begin() as conn:
+ conn.execute(
+ Query.__table__.update()
+ .where(Query.id == 123)
+ .values(status=QueryStatusEnum.STOPPED, end_time=1234.5)
+ )
+ query.rows = 42
+ query.progress = 100
+ query.end_time = 5678.9
+ query.set_extra_json_key("columns", [{"name": "id"}])
+
+ assert (
+ celery_task._make_check_stopped_fn(query, preserve_pending=True)()
+ == stopped
+ )
+ if stopped:
+ payload = celery_task._stopped_query_payload(query)
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ assert query.end_time == 1234.5
+ assert query.rows != 42
+ assert "columns" not in query.extra
+ else:
+ query.status = QueryStatusEnum.SUCCESS
+ assert celery_task._commit_query_unless_stopped(query)
+ session.refresh(query)
+ assert query.status == QueryStatusEnum.SUCCESS
+ assert query.rows == 42
+ assert query.progress == 100
+ assert query.end_time == Decimal("5678.9")
+ assert query.extra["columns"] == [{"name": "id"}]
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ "checkpoint", ["serialization", "backend_error", "timeout", "timeout_oauth2_lookup"]
+)
+@pytest.mark.parametrize("stop_requested", [False, True])
+def test_task_error_lookup_does_not_flush_over_a_stop(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ checkpoint: str,
+ stop_requested: bool,
+) -> None:
+ """
+ The task's error lookup must not flush stale metadata before checking Stop,
+ including when a time-limit signal fires after the Stop.
+ """
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_DISPATCHED_KEY, QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'task-stop.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+
+ def interrupt_processing() -> None:
+ """Optionally commit Stop independently before result processing fails."""
+ if stop_requested:
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+ if checkpoint.startswith("timeout"):
+ raise SoftTimeLimitExceeded()
+ raise RuntimeError("Result processing failed")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="task-stop",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+ execute = mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ mocker.patch.object(celery_task, "results_backend", None)
+ if checkpoint.startswith("timeout"):
+ execute.side_effect = lambda **kwargs: interrupt_processing()
+ if checkpoint == "timeout_oauth2_lookup":
+
+ def load_oauth2_settings() -> bool:
+ """Model the ORM load while deciding whether to retry an error."""
+ session.query(Query).count()
+ return False
+
+ mock_database.is_oauth2_enabled.side_effect = load_oauth2_settings
+ elif checkpoint == "serialization":
+ result_set = MagicMock(size=1, truncated=False)
+ execute.return_value = [("SELECT 1", result_set, 1.0, 1)]
+ mocker.patch.object(
+ celery_task,
+ "_serialize_result_set",
+ side_effect=lambda _: interrupt_processing(),
+ )
+ else:
+ backend = MagicMock()
+ backend.set.side_effect = lambda *args: interrupt_processing()
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+
+ # Keep the real _get_query: a mocked lookup cannot expose autoflush.
+ payload = celery_task.execute_sql_task(123, query.sql)
+
+ assert payload is not None
+ timed_out = checkpoint.startswith("timeout")
+ if stop_requested:
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ else:
+ expected_status = (
+ QueryStatusEnum.TIMED_OUT if timed_out else QueryStatusEnum.FAILED
+ )
+ assert payload["status"] == expected_status
+ session.refresh(query)
+ if stop_requested:
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.end_time == Decimal("1234.5")
+ assert query.extra == {
+ QUERY_EARLY_CANCEL_KEY: True,
+ QUERY_DISPATCHED_KEY: True,
+ }
+ assert query.error_message is None
+ else:
+ assert query.status == expected_status
+ assert query.end_time is not None
+ assert query.error_message is not None
+ assert query.results_key is None
+ engine.dispose()
+
+
+def test_timeout_is_guarded_against_a_stop(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_db_session: MagicMock,
+) -> None:
+ """A timeout checks and guards against a previously accepted Stop."""
+ from superset.sql.execution import celery_task
+
+ mock_query.status = QueryStatusEnum.TIMED_OUT
+ mock_query.tmp_table_name = "pending_table"
+
+ payload = celery_task._handle_query_error(RuntimeError("Timed out"), mock_query)
+
+ assert payload["status"] == QueryStatusEnum.TIMED_OUT
+ assert mock_query.status == QueryStatusEnum.TIMED_OUT
+ mock_db_session.refresh.assert_called_once()
+ mock_db_session.query.assert_called_once()
+ mock_db_session.flush.assert_not_called()
+ mock_db_session.commit.assert_called_once()
+
+
+@pytest.mark.parametrize("mutation", [False, True], ids=["select", "mutation"])
+@pytest.mark.parametrize(
+ "checkpoint",
+ ["success", "backend_failure", "error", "payload_load", "error_load"],
+)
+def test_stop_cannot_be_overwritten_by_terminal_metadata(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ checkpoint: str,
+ mutation: bool,
+) -> None:
+ """A separately committed stop must win even after the worker's last read."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'stop.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ original_make_check = celery_task._make_check_stopped_fn
+ stop_delivered = False
+
+ def deliver_stop() -> None:
+ """Commit the stop in an independent transaction."""
+ nonlocal stop_delivered
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+ stop_delivered = True
+
+ def make_check(query: Query, preserve_pending: bool = False) -> Callable[[], bool]:
+ """Deliver Stop after a successful read, before any terminal writes."""
+ check = original_make_check(query, preserve_pending)
+
+ def check_then_stop() -> bool:
+ """Expose the check/write race without mocking ORM operations."""
+ stopped = check()
+ if (
+ preserve_pending
+ and not stop_delivered
+ and checkpoint not in {"payload_load", "error_load"}
+ ):
+ deliver_stop()
+ return stopped
+
+ return check_then_stop
+
+ mocker.patch.object(celery_task, "_make_check_stopped_fn", side_effect=make_check)
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="stop-race",
+ database_id=1,
+ sql="INSERT INTO users VALUES (1)" if mutation else "SELECT 1",
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ session.refresh(query)
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ mocker.patch.object(celery_task, "results_backend", None)
+ if checkpoint == "backend_failure":
+ backend = MagicMock()
+ backend.set.return_value = False
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+
+ def load_relationship() -> Any:
+ """Model the autoflush caused by a lazy relationship lookup."""
+ deliver_stop()
+ session.query(Query).count()
+ return {"id": 123} if checkpoint == "payload_load" else []
+
+ if checkpoint == "payload_load":
+ mocker.patch.object(Query, "to_dict", side_effect=load_relationship)
+ elif checkpoint == "error_load":
+ mock_database.db_engine_spec.extract_errors.side_effect = (
+ lambda *args, **kwargs: load_relationship()
+ )
+
+ payload: dict[str, Any] | None
+ if checkpoint in {"error", "error_load"}:
+ payload = celery_task._handle_query_error(
+ RuntimeError("Driver error"), query
+ )
+ else:
+ payload = celery_task._execute_sql_statements(123, query.sql, None)
+
+ assert stop_delivered
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ session.refresh(query)
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.end_time == Decimal("1234.5")
+ assert query.extra[QUERY_EARLY_CANCEL_KEY] is True
+ assert query.error_message is None
+ assert query.results_key is None
+ engine.dispose()
+
+
+@pytest.mark.parametrize("backend_write", [None, True, False])
+def test_stop_during_mutation_commit_discards_results(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ backend_write: bool | None,
+) -> None:
+ """An accepted stop keeps STOPPED even if the mutation committed."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'commit.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ def deliver_stop() -> None:
+ """Commit the stop after the pre-commit check, mid target commit."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ conn = setup_mock_raw_connection(mock_database)
+ conn.commit.side_effect = deliver_stop
+ sql = "INSERT INTO users VALUES (1)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="commit-race",
+ database_id=1,
+ sql=sql,
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ session.refresh(query)
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task, "execute_sql_with_cursor", return_value=[(sql, None, 1.0, 1)]
+ )
+ backend = None
+ if backend_write is not None:
+ backend = MagicMock()
+ backend.set.return_value = backend_write
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"x")
+ mocker.patch.object(celery_task, "results_backend", backend)
+
+ payload = celery_task._execute_sql_statements(123, sql, None)
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ conn.commit.assert_called_once_with()
+ session.refresh(query)
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.results_key is None
+ if backend:
+ backend.set.assert_not_called()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [False, True], ids=["mutation", "ctas"])
+@pytest.mark.parametrize("stop_intent", ["status", "early"])
+@pytest.mark.parametrize("seen", [False, True], ids=["after_finish", "while_running"])
+@pytest.mark.parametrize("autocommit", [True, False], ids=["autocommit", "transaction"])
+def test_autocommit_work_finished_before_a_stop_is_committed(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ autocommit: bool,
+ seen: bool,
+ stop_intent: str,
+ ctas: bool,
+) -> None:
+ """An accepted stop discards autocommit results even if not locally seen."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'autocommit.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ conn = setup_mock_raw_connection(mock_database)
+ conn.dbapi_connection.autocommit = autocommit
+ sql = "SELECT 1" if ctas else "INSERT INTO users VALUES (1)"
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="autocommit",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+
+ def execute_then_stop(**kwargs: Any) -> list[tuple[str, None, float, int]]:
+ """Finish the last statement, with a stop committed elsewhere."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ if stop_intent == "early":
+ stopped_query.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ else:
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+ if seen:
+ # An engine polling the running statement loads the stop.
+ session.refresh(query)
+ return [(sql, None, 1.0, 1)]
+
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task, "execute_sql_with_cursor", side_effect=execute_then_stop
+ )
+ mocker.patch.object(celery_task, "results_backend", None)
+
+ payload = celery_task._execute_sql_statements(123, sql, None)
+
+ expected = QueryStatusEnum.STOPPED
+ assert payload is not None
+ assert payload["status"] == expected
+ session.refresh(query)
+ assert query.status == expected
+ assert not conn.commit.called
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [True, False], ids=["ctas", "insert"])
+def test_sqlite_ctas_finished_before_a_stop_is_committed(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ ctas: bool,
+) -> None:
+ """A durable SQLite CTAS remains stopped without undoing its table."""
+ from contextlib import contextmanager
+
+ from sqlalchemy import create_engine, inspect
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}")
+ Query.__table__.create(engine)
+ target = create_engine(f"sqlite:///{tmp_path / 'target.db'}")
+ with target.begin() as setup:
+ setup.exec_driver_sql("CREATE TABLE existing (a INTEGER)")
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ connection = target.raw_connection()
+ try:
+ yield connection
+ finally:
+ connection.close()
+
+ mock_database.get_raw_connection = raw_connection
+ sql = "SELECT 1 AS a" if ctas else "INSERT INTO existing VALUES (1)"
+ executed = "CREATE TABLE tmp_table AS SELECT 1 AS a" if ctas else sql
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="sqlite-ctas",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+
+ def execute_then_stop(
+ *, cursor: Any, **kwargs: Any
+ ) -> list[tuple[str, None, float, int]]:
+ """Run the statement on SQLite, then commit a stop elsewhere."""
+ cursor.execute(executed)
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stop_session.commit()
+ return [(executed, None, 1.0, 1)]
+
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task, "execute_sql_with_cursor", side_effect=execute_then_stop
+ )
+ mocker.patch.object(celery_task, "results_backend", None)
+
+ payload = celery_task._execute_sql_statements(123, sql, None)
+
+ expected = QueryStatusEnum.STOPPED
+ assert payload is not None
+ assert payload["status"] == expected
+ session.refresh(query)
+ assert query.status == expected
+ with target.connect() as observer:
+ assert inspect(observer).has_table("tmp_table") is ctas
+ rows = observer.exec_driver_sql("SELECT COUNT(*) FROM existing").scalar()
+ assert rows == 0
+ target.dispose()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [False, True], ids=["mutation", "ctas"])
+@pytest.mark.parametrize(
+ "failure",
+ [
+ "backend_set_raises",
+ "serialize_raises",
+ "payload_too_large",
+ "backend_write_false",
+ "connection_close_raises",
+ ],
+)
+def test_failure_after_stop_raced_commit_preserves_stopped(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ ctas: bool,
+ failure: str,
+) -> None:
+ """Even a connection-close failure cannot overwrite an accepted stop."""
+ from contextlib import contextmanager
+
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import create_mock_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'commit.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mock_database.is_oauth2_enabled.return_value = False
+
+ def deliver_stop() -> None:
+ """Commit the stop after the pre-commit check, mid target commit."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatusEnum.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ conn = create_mock_connection()
+ conn.commit.side_effect = deliver_stop
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ yield conn
+ if failure == "connection_close_raises":
+ raise RuntimeError("Connection close failed")
+
+ mock_database.get_raw_connection = raw_connection
+
+ sql = "SELECT 1" if ctas else "INSERT INTO users VALUES (1)"
+ backend = MagicMock()
+ backend.set.return_value = failure != "backend_write_false"
+ if failure == "backend_set_raises":
+ backend.set.side_effect = RuntimeError("Backend unavailable")
+ serialize = mocker.patch.object(celery_task, "_serialize_payload")
+ serialize.return_value = b"x" * (2 * 1024 * 1024)
+ if failure == "serialize_raises":
+ serialize.side_effect = RuntimeError("Cannot serialize")
+ mocker.patch.object(celery_task, "results_backend", backend)
+ mocker.patch.object(celery_task, "security_manager")
+ mocker.patch.dict(
+ current_app.config,
+ {
+ "STATS_LOGGER": MagicMock(),
+ "SQLLAB_PAYLOAD_MAX_MB": 1 if failure == "payload_too_large" else None,
+ },
+ )
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="commit-race-failure",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ status=QueryStatusEnum.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ session.refresh(query)
+ mocker.patch.object(celery_task.db, "session", session)
+ mocker.patch.object(celery_task, "_get_query", return_value=query)
+ mocker.patch.object(
+ celery_task, "execute_sql_with_cursor", return_value=[(sql, None, 1.0, 1)]
+ )
+
+ payload = celery_task.execute_sql_task(123, sql)
+
+ conn.commit.assert_called_once_with()
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ session.refresh(query)
+ assert query.status == QueryStatusEnum.STOPPED
+ assert query.error_message is None
+ assert query.results_key is None
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ "status",
+ [QueryStatusEnum.RUNNING, QueryStatusEnum.FAILED],
+ ids=["success", "failed"],
+)
+def test_publish_query_result_preserves_failure(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ status: QueryStatusEnum,
+) -> None:
+ """
+ Publishing keeps a committed backend failure for reads and committed
+ mutations, and otherwise publishes SUCCESS.
+ """
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'publish.db'}")
+ Query.__table__.create(engine)
+ with Session(engine) as session:
+ query = Query(
+ id=123, client_id="publish", database_id=1, sql="SELECT 1", status=status
+ )
+ session.add(query)
+ # _store_results_in_backend commits FAILED before publishing can run.
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+ query.rows = 1
+ payload = {"status": QueryStatusEnum.SUCCESS.value, "query_id": 123}
+
+ assert celery_task._publish_query_result(query, payload) is payload
+
+ session.expire_all()
+ assert query.status == (
+ QueryStatusEnum.FAILED
+ if status == QueryStatusEnum.FAILED
+ else QueryStatusEnum.SUCCESS
+ )
+ assert query.rows == 1
+
+
+def test_successful_select_reads_stop_state_once_per_checkpoint(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+) -> None:
+ """
+ A successful SELECT reads stop state before results are serialized, after
+ the results backend write and in the guarded publish, not again in between.
+ """
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'task-reads.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+ result_set = MagicMock(size=1, truncated=False, columns=[])
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", result_set, 1.0, 1)],
+ )
+ mocker.patch.object(celery_task, "_serialize_result_set", return_value=([], []))
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ backend = MagicMock()
+ backend.set.return_value = True
+ mocker.patch.object(celery_task, "results_backend", backend)
+
+ with Session(engine) as session:
+ session.add(
+ Query(
+ id=123,
+ client_id="task-reads",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.PENDING,
+ )
+ )
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+ statements: list[str] = []
+ event.listen(
+ engine,
+ "before_cursor_execute",
+ lambda conn, cursor, statement, *args: statements.append(statement),
+ )
+
+ payload = celery_task.execute_sql_task(123, "SELECT 1")
+
+ assert payload is not None
+ assert payload["status"] == QueryStatusEnum.SUCCESS
+ full_refreshes = [
+ sql
+ for sql in statements
+ if sql.startswith('SELECT "query".id, "query".client_id')
+ ]
+ status_reads = [
+ sql for sql in statements if sql.startswith('SELECT "query".status')
+ ]
+ # One refresh before the pre-commit decision; status-only reads after
+ # the backend write and before the guarded publish.
+ assert len(full_refreshes) == 1
+ assert len(status_reads) == 2
+
+
+@pytest.mark.parametrize("extra_json", [None, b"{}", MagicMock(), "", "{}"])
+@pytest.mark.parametrize("mapped", [False, True])
+def test_stop_seen_during_execution_ignores_non_text_extra_json(
+ extra_json: Any,
+ mapped: bool,
+) -> None:
+ """Only loaded text JSON is decoded, including when the query is a double."""
+ from superset.models.sql_lab import Query
+ from superset.sql.execution.celery_task import stop_seen_during_execution
+
+ query_type = Query if mapped else MagicMock
+ query = query_type(status=QueryStatusEnum.RUNNING, extra_json=extra_json)
+ assert not stop_seen_during_execution(query)
+
+
+@pytest.mark.parametrize("intent", ["status", "early", "none"])
+@pytest.mark.parametrize("expired", [False, True])
+def test_stop_seen_during_execution_reads_only_loaded_attributes(
+ intent: str, expired: bool, mocker: MockerFixture
+) -> None:
+ """Expired stop attributes stay unseen without triggering a database read."""
+ from sqlalchemy import create_engine, event
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.sql.execution.celery_task import stop_seen_during_execution
+
+ engine = create_engine("sqlite://")
+ Query.__table__.create(engine)
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="loaded-stop",
+ database_id=1,
+ sql="SELECT 1",
+ status=(
+ QueryStatusEnum.STOPPED
+ if intent == "status"
+ else QueryStatusEnum.RUNNING
+ ),
+ extra_json=superset_json.dumps({QUERY_EARLY_CANCEL_KEY: intent == "early"}),
+ )
+ session.add(query)
+ session.flush()
+ if expired:
+ session.expire(query, ["status", "extra_json"])
+ execute = mocker.Mock()
+ event.listen(engine, "before_cursor_execute", execute)
+ assert stop_seen_during_execution(query) is (not expired and intent != "none")
+ execute.assert_not_called()
+ engine.dispose()
+
+
+def test_replayed_task_stop_keeps_results_the_committed_row_references(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+) -> None:
+ """A redelivered task that a stop wins keeps the earlier attempt's results."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'replay.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(celery_task.security_manager, "find_user", return_value=None)
+ mock_database.is_oauth2_enabled.return_value = False
+ mocker.patch.object(
+ celery_task,
+ "execute_sql_with_cursor",
+ return_value=[("SELECT 1", None, 1.0, 0)],
+ )
+ stored = {"prior-attempt-key": b"old data"}
+
+ def store_then_stop(key: str, value: bytes, timeout: int) -> bool:
+ stored[key] = value
+ with Session(engine) as stopper:
+ stopped = stopper.get(Query, 123)
+ stopped.status = QueryStatusEnum.STOPPED
+ stopper.commit()
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store_then_stop
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(celery_task, "results_backend", backend)
+ with Session(engine) as session:
+ session.add(
+ Query(
+ id=123,
+ client_id="replayed",
+ database_id=1,
+ sql="SELECT 1",
+ status=QueryStatusEnum.SUCCESS,
+ results_key="prior-attempt-key",
+ )
+ )
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ payload = celery_task.execute_sql_task(123, "SELECT 1")
+ assert payload is not None
+
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ with Session(engine) as observer:
+ row = observer.get(Query, 123)
+ assert row.status == QueryStatusEnum.STOPPED
+ assert row.results_key == "prior-attempt-key"
+ # Only this attempt's unpublished blob is deleted.
+ assert list(stored) == ["prior-attempt-key"]
+ assert backend.delete.call_count == 1
+ engine.dispose()
+
+
+@pytest.mark.parametrize("autocommit", [False, True], ids=["transaction", "autocommit"])
+@pytest.mark.parametrize("intent", ["status", "early"])
+def test_stop_seen_by_engine_survives_progress_commit(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ autocommit: bool,
+ intent: str,
+) -> None:
+ """
+ An engine that cancels a running mutation on a stop ends STOPPED, even
+ though the shared executor's progress commit expires the loaded stop.
+ """
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'cancel-progress.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task, "results_backend", None)
+ conn = setup_mock_raw_connection(mock_database)
+ cursor = conn.cursor.return_value
+ cursor.description = None
+ cursor.rowcount = 0
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="cancel-progress",
+ database_id=1,
+ sql="INSERT INTO users VALUES (1)",
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ def poll_and_cancel(cursor: Any, sql: str, observed: Query) -> None:
+ with Session(engine) as stopper:
+ stopped = stopper.get(Query, 123)
+ if intent == "status":
+ stopped.status = QueryStatusEnum.STOPPED
+ else:
+ stopped.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ stopper.commit()
+ # A polling engine loads the stop, cancels the operation and
+ # returns with no cursor description.
+ session.refresh(observed)
+ cursor.cancel()
+
+ mock_database.db_engine_spec.execute_with_cursor.side_effect = poll_and_cancel
+
+ payload = celery_task._execute_sql_statements(123, query.sql, None)
+ assert payload is not None
+
+ cursor.cancel.assert_called_once()
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ with Session(engine) as observer:
+ assert observer.get(Query, 123).status == QueryStatusEnum.STOPPED
+ conn.commit.assert_not_called()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("intent", ["status", "early"])
+def test_impala_cancel_survives_progress_commit(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+ intent: str,
+) -> None:
+ """Impala's own polling cancel of an autocommit INSERT ends STOPPED."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.constants import QUERY_EARLY_CANCEL_KEY
+ from superset.db_engine_specs.impala import ImpalaEngineSpec
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'impala-cancel.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task, "results_backend", None)
+ mock_database.db_engine_spec = ImpalaEngineSpec
+ conn = setup_mock_raw_connection(mock_database)
+ # Only the RPC cursor is a double: execute_with_cursor, handle_cursor and
+ # cancellation are Impala's own.
+ cursor = conn.cursor.return_value
+ cursor.description = None
+ cursor.rowcount = 0
+ cursor.status.return_value = "RUNNING_STATE"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="impala-cancel",
+ database_id=1,
+ sql="INSERT INTO users VALUES (1)",
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ def start_then_stop(sql: str) -> None:
+ with Session(engine) as stopper:
+ stopped = stopper.get(Query, 123)
+ if intent == "status":
+ stopped.status = QueryStatusEnum.STOPPED
+ else:
+ stopped.set_extra_json_key(QUERY_EARLY_CANCEL_KEY, True)
+ stopper.commit()
+
+ cursor.execute_async.side_effect = start_then_stop
+
+ payload = celery_task._execute_sql_statements(123, query.sql, None)
+ assert payload is not None
+
+ cursor.cancel_operation.assert_called_once()
+ cursor.close_operation.assert_called_once()
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ with Session(engine) as observer:
+ assert observer.get(Query, 123).status == QueryStatusEnum.STOPPED
+ conn.commit.assert_not_called()
+ engine.dispose()
+
+
+def test_stop_seen_between_autocommit_statements_ends_stopped(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_database: MagicMock,
+ tmp_path: Path,
+) -> None:
+ """A stop the check before the next statement loads ends the script STOPPED."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'between-statements.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(celery_task, "results_backend", None)
+ mock_database.db_engine_spec.run_multiple_statements_as_one = False
+ conn = setup_mock_raw_connection(mock_database)
+ cursor = conn.cursor.return_value
+ cursor.description = None
+ cursor.rowcount = 1
+ sql = "INSERT INTO users VALUES (1);\nINSERT INTO users VALUES (2)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="between-statements",
+ database_id=1,
+ sql=sql,
+ status=QueryStatusEnum.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(celery_task.db, "session", session)
+
+ def finish_then_stop(cursor: Any, sql: str, observed: Query) -> None:
+ # The first statement finishes; a stop lands before the next one.
+ with Session(engine) as stopper:
+ stopper.get(Query, 123).status = QueryStatusEnum.STOPPED
+ stopper.commit()
+
+ execute = mock_database.db_engine_spec.execute_with_cursor
+ execute.side_effect = finish_then_stop
+
+ payload = celery_task._execute_sql_statements(123, sql, None)
+ assert payload is not None
+
+ execute.assert_called_once()
+ assert payload["status"] == QueryStatusEnum.STOPPED
+ with Session(engine) as observer:
+ assert observer.get(Query, 123).status == QueryStatusEnum.STOPPED
+ conn.commit.assert_not_called()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("stop_before_read", [False, True])
+def test_celery_startup_preserves_stop_before_running_write(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Path,
+ mock_database: MagicMock,
+ stop_before_read: bool,
+) -> None:
+ """Celery startup must not replace a committed Stop with RUNNING."""
+ from sqlalchemy import create_engine
+ from sqlalchemy.orm import Session
+
+ from superset.models.sql_lab import Query
+ from superset.sql.execution import celery_task
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'celery-startup.db'}")
+ Query.__table__.create(engine)
+ mocker.patch.object(Query, "database", PropertyMock(return_value=mock_database))
+
+ def stop() -> None:
+ """Commit after retrieval but before the startup status write."""
+ with Session(engine) as stop_session:
+ stopped = stop_session.get(Query, 123)
+ assert stopped is not None
+ stopped.status = QueryStatusEnum.STOPPED
+ stopped.end_time = 1234.5
+ stop_session.commit()
+
+ if not stop_before_read:
+ mock_database.db_engine_spec.patch.side_effect = stop
+ parse = mocker.patch.object(
+ celery_task,
+ "SQLScript",
+ side_effect=AssertionError("Stopped query reached parsing"),
+ )
+ with Session(engine) as worker_session:
+ worker_session.add(
+ Query(
+ id=123,
+ client_id="startup",
+ database_id=1,
+ status=QueryStatusEnum.PENDING,
+ )
+ )
+ worker_session.commit()
+ if stop_before_read:
+ stop()
+ mocker.patch.object(celery_task.db, "session", worker_session)
+ payload = celery_task._execute_sql_statements(123, "SELECT 1", None)
+ assert payload == {"query_id": 123, "status": QueryStatusEnum.STOPPED}
+ parse.assert_not_called()
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatusEnum.STOPPED
+ assert persisted.end_time == 1234.5
+ engine.dispose()
+
+
+@pytest.mark.parametrize("cancel_id", [None, "cancel-123"])
+def test_celery_records_dispatch_before_execution_without_cancel_handle(
+ mocker: MockerFixture,
+ app_context: None,
+ mock_query: MagicMock,
+ mock_database: MagicMock,
+ cancel_id: str | None,
+) -> None:
+ """A dispatched connector without a cancel ID must not claim early cancel."""
+ from superset.constants import QUERY_DISPATCHED_KEY
+ from superset.sql.execution import celery_task
+
+ from .conftest import setup_mock_raw_connection
+
+ mock_query.database = mock_database
+ mock_database.db_engine_spec.get_cancel_query_id.return_value = cancel_id
+ setup_mock_raw_connection(mock_database)
+ mocker.patch.object(celery_task, "_get_query", return_value=mock_query)
+
+ def execute(**kwargs: Any) -> list[Any]:
+ """Check dispatch metadata at the executor boundary."""
+ mock_query.set_extra_json_key.assert_any_call(QUERY_DISPATCHED_KEY, True)
+ return []
+
+ mocker.patch.object(celery_task, "execute_sql_with_cursor", side_effect=execute)
+ celery_task._execute_sql_statements(123, "SELECT 1", None)
diff --git a/tests/unit_tests/sql/execution/test_executor.py b/tests/unit_tests/sql/execution/test_executor.py
index 329342430a6..cca0ae06735 100644
--- a/tests/unit_tests/sql/execution/test_executor.py
+++ b/tests/unit_tests/sql/execution/test_executor.py
@@ -978,8 +978,7 @@ def test_async_handle_cancel(
from superset.models.sql_lab import Query
from superset.sql.execution.executor import SQLExecutor
- mock_query = MagicMock(spec=Query)
- mock_query.status = "running"
+ mock_query = Query(status="running")
mock_query.database = database
filter_mock = mock_db_session.query.return_value.filter_by.return_value
filter_mock.one_or_none.return_value = mock_query
diff --git a/tests/unit_tests/sql/execution/test_utils.py b/tests/unit_tests/sql/execution/test_utils.py
new file mode 100644
index 00000000000..6233568ce24
--- /dev/null
+++ b/tests/unit_tests/sql/execution/test_utils.py
@@ -0,0 +1,207 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for helpers shared by legacy and Celery SQL execution."""
+
+import inspect
+
+import pytest
+from flask import current_app
+from pytest_mock import MockerFixture
+from sqlalchemy.orm import Session
+
+from superset.common.db_query_status import QueryStatus
+from superset.constants import QUERY_TARGET_COMMITTED_KEY
+from superset.models.sql_lab import Query
+from superset.sql.execution.utils import (
+ commit_target_transaction,
+ finish_query_stop,
+ set_query_stopped,
+)
+
+
+@pytest.mark.parametrize(
+ ("legacy_name", "celery_name"),
+ [
+ ("_extract_errors", "_extract_errors"),
+ ("_discard_unreferenced_results", "_discard_unreferenced_results"),
+ ("_discard_unpublished_results", "_discard_stored_results"),
+ ],
+)
+def test_execution_helpers_are_shared(legacy_name: str, celery_name: str) -> None:
+ """Both execution paths use the same neutral helper implementation."""
+ from superset import sql_lab
+ from superset.sql.execution import celery_task
+
+ helper = getattr(sql_lab, legacy_name)
+ assert helper is getattr(celery_task, celery_name)
+ assert helper.__module__ == "superset.sql.execution.utils"
+
+
+@pytest.fixture
+def persisted_stop_query(session: Session) -> Query:
+ """Create a query on the isolated session used by the production helpers."""
+ Query.__table__.create(session.get_bind())
+ query = Query(
+ id=123,
+ client_id="helper-stop",
+ database_id=1,
+ user_id=7,
+ status=QueryStatus.RUNNING,
+ )
+ session.add(query)
+ session.commit()
+ return query
+
+
+@pytest.mark.parametrize(
+ "status", [QueryStatus.SUCCESS, QueryStatus.FAILED, QueryStatus.TIMED_OUT]
+)
+def test_stop_rejects_completed_query(
+ session: Session, persisted_stop_query: Query, status: QueryStatus
+) -> None:
+ """A losing Stop leaves the completed outcome and timestamp untouched."""
+ query = persisted_stop_query
+ query.status = status
+ query.end_time = 42.0
+ session.commit()
+
+ assert set_query_stopped(query, 84.0) is False
+ session.commit()
+ session.refresh(query)
+ assert query.status == status
+ assert query.end_time == 42.0
+
+
+def test_target_commit_skips_accepted_stop(
+ mocker: MockerFixture, session: Session, persisted_stop_query: Query
+) -> None:
+ """An accepted Stop prevents the supplied target commit from running."""
+ query = persisted_stop_query
+ query.status = QueryStatus.STOPPED
+ query.end_time = 42.0
+ session.commit()
+ commit = mocker.Mock()
+
+ assert commit_target_transaction(query, commit) is False
+ commit.assert_not_called()
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time == 42.0
+ assert QUERY_TARGET_COMMITTED_KEY not in query.extra
+
+
+@pytest.mark.parametrize("cancel_outcome", ["success", "unsupported", "unreachable"])
+@pytest.mark.parametrize("metadata_failure", ["none", "merge", "commit"])
+def test_stop_acceptance_survives_cancellation_failures(
+ mocker: MockerFixture,
+ session: Session,
+ persisted_stop_query: Query,
+ cancel_outcome: str,
+ metadata_failure: str,
+) -> None:
+ """Failure after the durable Stop cannot reject it or repeat cancellation."""
+ query = persisted_stop_query
+ assert set_query_stopped(query, 42.0) is True
+ cancel = mocker.Mock(return_value=cancel_outcome == "success")
+ if cancel_outcome == "unreachable":
+ cancel.side_effect = RuntimeError("Engine unavailable")
+ if metadata_failure == "merge":
+ mocker.patch(
+ "superset.sql.execution.utils.set_query_stopped",
+ side_effect=RuntimeError("Cancellation metadata unavailable"),
+ )
+ commit = session.commit
+ commits = 0
+
+ def fail_second_commit() -> None:
+ """Fail only after the initial Stop is durably accepted."""
+ nonlocal commits
+ commits += 1
+ if metadata_failure == "commit" and commits == 2:
+ raise RuntimeError("Cancellation metadata commit unavailable")
+ commit()
+
+ mocker.patch.object(session, "commit", side_effect=fail_second_commit)
+
+ assert finish_query_stop(query, cancel) is True
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time == 42.0
+ session.expire_all()
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time == 42.0
+ assert finish_query_stop(query, cancel) is True
+ cancel.assert_called_once_with()
+
+
+def test_async_cancel_rejects_target_commit_winner(
+ mocker: MockerFixture, session: Session, persisted_stop_query: Query
+) -> None:
+ """An async handle must propagate a target commit that defeats its Stop."""
+ from superset.sql.execution.executor import SQLExecutor
+
+ query = persisted_stop_query
+ query.set_extra_json_key(QUERY_TARGET_COMMITTED_KEY, True)
+ session.commit()
+ cancel = mocker.patch.object(SQLExecutor, "_cancel_query")
+
+ assert SQLExecutor(mocker.Mock())._create_async_handle(query.id).cancel() is False
+ cancel.assert_not_called()
+ session.refresh(query)
+ assert query.status == QueryStatus.RUNNING
+ assert query.extra[QUERY_TARGET_COMMITTED_KEY] is True
+ assert query.end_time is None
+
+
+@pytest.mark.parametrize("caller", ["helper", "api", "async_handle"])
+def test_initial_stop_commit_failure_does_not_accept_or_cancel(
+ mocker: MockerFixture,
+ app_context: None,
+ session: Session,
+ persisted_stop_query: Query,
+ caller: str,
+) -> None:
+ """Only a durable initial commit may report acceptance or cancel the engine."""
+ from superset.queries.api import QueryRestApi
+ from superset.sql.execution.executor import SQLExecutor
+
+ query = persisted_stop_query
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ cancel = mocker.Mock()
+ mocker.patch("superset.sql_lab.cancel_query", cancel)
+ mocker.patch.object(SQLExecutor, "_cancel_query", cancel)
+ mocker.patch.object(
+ session, "commit", side_effect=RuntimeError("Initial Stop commit failed")
+ )
+
+ def request_stop() -> None:
+ """Exercise the selected acceptance entry point without swallowing errors."""
+ if caller == "helper":
+ assert set_query_stopped(query, 42.0) is True
+ finish_query_stop(query, cancel)
+ elif caller == "api":
+ with current_app.test_request_context(json={"client_id": query.client_id}):
+ inspect.unwrap(QueryRestApi.stop_query)(QueryRestApi())
+ else:
+ SQLExecutor(mocker.Mock())._create_async_handle(query.id).cancel()
+
+ with pytest.raises(RuntimeError, match="Initial Stop commit failed"):
+ request_stop()
+
+ cancel.assert_not_called()
+ session.rollback()
+ session.refresh(query)
+ assert query.status == QueryStatus.RUNNING
+ assert query.end_time is None
diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py
index 77817152bf4..aa9b04400a9 100644
--- a/tests/unit_tests/sql_lab_test.py
+++ b/tests/unit_tests/sql_lab_test.py
@@ -221,6 +221,7 @@ def test_execute_sql_statement_exceeds_payload_limit(
# Mock db.session.refresh to avoid AttributeError during session refresh
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
# Mock the results backend to avoid "Results backend is not configured" error
mocker.patch("superset.sql_lab.results_backend", return_value=True)
@@ -280,6 +281,7 @@ def test_execute_sql_statement_within_payload_limit(mocker: MockerFixture, app)
# Mock db.session.refresh to avoid AttributeError during session refresh
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
# Mock the results backend to avoid "Results backend is not configured" error
mocker.patch("superset.sql_lab.results_backend", return_value=True)
@@ -321,6 +323,7 @@ def test_execute_sql_statements_rejects_client_file_transfer(
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
with pytest.raises(SupersetDisallowedClientFileTransferException) as excinfo:
execute_sql_statements(
@@ -373,6 +376,7 @@ def test_execute_sql_statements_mutates_before_split_by_default(
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
@@ -436,6 +440,7 @@ def test_execute_sql_statements_mutates_per_statement_when_run_as_one(
side_effect=lambda payload, use_msgpack: "serialized_payload",
)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
execute_sql_statements(
@@ -491,6 +496,7 @@ def test_execute_sql_statements_raises_when_mutator_strips_all_statements(
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
@@ -535,6 +541,7 @@ def test_execute_sql_statements_raises_when_mutator_strips_single_block(
mocker.patch("superset.sql_lab.get_query", return_value=query)
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
mocker.patch("superset.sql_lab.results_backend", return_value=True)
with pytest.raises(SupersetErrorException):
@@ -577,6 +584,7 @@ def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
# FAILED; `query` here is a MagicMock, not a real persistent ORM
# instance, so the real refresh() would error introspecting it.
mocker.patch("superset.sql_lab.db.session.refresh", return_value=None)
+ mocker.patch("superset.sql_lab.commit_query_unless_stopped", return_value=True)
g = mocker.patch("superset.db_engine_specs.base.g")
g.user = mocker.MagicMock()
@@ -648,6 +656,1325 @@ def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
app_context.pop()
+@pytest.mark.parametrize("ctas", [False, True], ids=["mutation", "ctas"])
+@pytest.mark.parametrize(
+ "failure",
+ [
+ None,
+ "backend_write_false",
+ "backend_set_raises",
+ "serialize_raises",
+ "payload_too_large",
+ "returned_payload_too_large",
+ "connection_close_raises",
+ ],
+)
+def test_get_sql_results_preserves_stop_when_mutation_commits(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ ctas: bool,
+ failure: str | None,
+) -> None:
+ """Test that a stop stays accepted even when the mutation commits."""
+ from contextlib import contextmanager
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'commit.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_dml = True
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.select_star.return_value = "SELECT * FROM tmp_table"
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ def deliver_stop() -> None:
+ """Commit the stop after the last per-block check, mid target commit."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ conn = MagicMock()
+ conn.commit.side_effect = deliver_stop
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ yield conn
+ if failure == "connection_close_raises":
+ raise RuntimeError("Connection close failed")
+
+ database.get_raw_connection = raw_connection
+
+ store_results = failure in {
+ "backend_write_false",
+ "backend_set_raises",
+ "serialize_raises",
+ "payload_too_large",
+ }
+ backend = MagicMock()
+ backend.set.return_value = failure != "backend_write_false"
+ if failure == "backend_set_raises":
+ backend.set.side_effect = RuntimeError("Backend unavailable")
+ mocker.patch.object(sql_lab, "results_backend", backend)
+ serialize = mocker.patch.object(sql_lab, "_serialize_payload")
+ serialize.return_value = b"x" * (2 * 1024 * 1024)
+ if failure == "serialize_raises":
+ serialize.side_effect = RuntimeError("Cannot serialize")
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ result_set = MagicMock(size=1, columns=[])
+ mocker.patch.object(sql_lab, "execute_query", return_value=result_set)
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.dict(
+ app.config,
+ {
+ "STATS_LOGGER": MagicMock(),
+ "SQLLAB_PAYLOAD_MAX_MB": (
+ 1
+ if failure in {"payload_too_large", "returned_payload_too_large"}
+ else None
+ ),
+ },
+ )
+
+ sql = "SELECT 1" if ctas else "INSERT INTO users VALUES (1)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="commit-race",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ ctas_method="TABLE",
+ tmp_table_name="tmp_table",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ payload = get_sql_results(
+ 123,
+ sql,
+ return_results=not store_results,
+ store_results=store_results,
+ )
+
+ conn.commit.assert_called_once_with()
+ expected = QueryStatus.STOPPED
+ assert payload is not None
+ assert payload["status"] == expected
+ session.refresh(query)
+ assert query.status == expected
+ assert query.results_key is None
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [False, True], ids=["mutation", "ctas"])
+def test_get_sql_results_stop_before_target_commit_stays_stopped(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ ctas: bool,
+) -> None:
+ """
+ Test that a stop committed after the last statement ran but before the
+ target database commit ends STOPPED: the mutation or CTAS is not committed
+ and the stop is not relabeled SUCCESS.
+ """
+ from contextlib import contextmanager
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'precommit.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_dml = True
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.select_star.return_value = "SELECT * FROM tmp_table"
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ conn = MagicMock()
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ yield conn
+
+ database.get_raw_connection = raw_connection
+
+ def execute_then_stop(*args: Any, **kwargs: Any) -> MagicMock:
+ """Finish the last statement, then commit a stop from another session."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_then_stop)
+ mocker.patch.object(sql_lab, "results_backend", None)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ sql = "SELECT 1" if ctas else "INSERT INTO users VALUES (1)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="precommit-stop",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ ctas_method="TABLE",
+ tmp_table_name="tmp_table",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ payload = sql_lab.execute_sql_statements(
+ 123,
+ sql,
+ True,
+ False,
+ None,
+ False,
+ None,
+ )
+
+ assert payload is not None
+ assert payload["status"] == QueryStatus.STOPPED
+ session.refresh(query)
+ assert query.status == QueryStatus.STOPPED
+ conn.commit.assert_not_called()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [False, True], ids=["mutation", "ctas"])
+@pytest.mark.parametrize("seen", [False, True], ids=["after_finish", "while_running"])
+@pytest.mark.parametrize("autocommit", [True, False], ids=["autocommit", "transaction"])
+def test_get_sql_results_autocommit_work_finished_before_a_stop_is_committed(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ autocommit: bool,
+ seen: bool,
+ ctas: bool,
+) -> None:
+ """Test that accepted stops discard autocommit results, even if unseen."""
+ from contextlib import contextmanager
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'autocommit.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_dml = True
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.select_star.return_value = "SELECT * FROM tmp_table"
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ conn = MagicMock()
+ conn.dbapi_connection.autocommit = autocommit
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ yield conn
+
+ database.get_raw_connection = raw_connection
+ mocker.patch.object(sql_lab, "results_backend", None)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.dict(app.config, {"STATS_LOGGER": MagicMock()})
+
+ sql = "SELECT 1" if ctas else "INSERT INTO users VALUES (1)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="autocommit-stop",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ ctas_method="TABLE",
+ tmp_table_name="tmp_table",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ def execute_then_stop(*args: Any, **kwargs: Any) -> MagicMock:
+ """Finish the last statement, with a stop committed elsewhere."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+ if seen:
+ # An engine polling the running statement loads the stop.
+ session.refresh(query)
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_then_stop)
+
+ payload = sql_lab.execute_sql_statements(
+ 123,
+ sql,
+ True,
+ False,
+ None,
+ False,
+ None,
+ )
+
+ expected = QueryStatus.STOPPED
+ assert payload is not None
+ assert payload["status"] == expected
+ session.refresh(query)
+ assert query.status == expected
+ assert not conn.commit.called
+ engine.dispose()
+
+
+@pytest.mark.parametrize("ctas", [True, False], ids=["ctas", "insert"])
+def test_get_sql_results_sqlite_ctas_finished_before_a_stop_is_committed(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ ctas: bool,
+) -> None:
+ """Test that a stopped SQLite CTAS can still leave a committed table."""
+ from contextlib import contextmanager
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, inspect
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'metadata.db'}")
+ Query.__table__.create(engine)
+ target = create_engine(f"sqlite:///{tmp_path / 'target.db'}")
+ with target.begin() as setup:
+ setup.exec_driver_sql("CREATE TABLE existing (a INTEGER)")
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_dml = True
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ database.select_star.return_value = "SELECT * FROM tmp_table"
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+
+ @contextmanager
+ def raw_connection(*args: Any, **kwargs: Any) -> Any:
+ connection = target.raw_connection()
+ try:
+ yield connection
+ finally:
+ connection.close()
+
+ database.get_raw_connection = raw_connection
+ mocker.patch.object(sql_lab, "results_backend", None)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.dict(app.config, {"STATS_LOGGER": MagicMock()})
+
+ sql = "SELECT 1 AS a" if ctas else "INSERT INTO existing VALUES (1)"
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="sqlite-ctas",
+ database_id=1,
+ sql=sql,
+ select_as_cta=ctas,
+ ctas_method="TABLE",
+ tmp_table_name="tmp_table",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ def execute_then_stop(query: Query, cursor: Any, log_params: Any) -> MagicMock:
+ """Run the statement on SQLite, then commit a stop elsewhere."""
+ cursor.execute(query.executed_sql)
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stop_session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_then_stop)
+
+ payload = sql_lab.execute_sql_statements(
+ 123, sql, True, False, None, False, None
+ )
+
+ expected = QueryStatus.STOPPED
+ assert payload is not None
+ assert payload["status"] == expected
+ session.refresh(query)
+ assert query.status == expected
+ with target.connect() as observer:
+ assert inspect(observer).has_table("tmp_table") is ctas
+ rows = observer.exec_driver_sql("SELECT COUNT(*) FROM existing").scalar()
+ assert rows == 0
+ target.dispose()
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ ("stop_at", "store_results", "return_results"),
+ [
+ ("serialize", False, True),
+ ("serialize", True, True),
+ ("serialize", True, False),
+ ("backend_write", True, True),
+ ("backend_write", True, False),
+ ],
+)
+def test_get_sql_results_stop_during_select_finalization_ends_stopped(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ stop_at: str,
+ store_results: bool,
+ return_results: bool,
+) -> None:
+ """
+ Test that a stop landing while a SELECT's results are serialized or stored
+ ends STOPPED in both the row and the returned payload, keeps the stop's end
+ time, and deletes stored results the stopped row does not reference.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'select-stop.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.dict(app.config, {"STATS_LOGGER": MagicMock(), "SQL_MAX_ROW": 0})
+
+ def deliver_stop() -> None:
+ """Commit the stop from another session."""
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stopped_query.end_time = 1234.5
+ stop_session.commit()
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ def serialize(*args: Any, **kwargs: Any) -> tuple[list[Any], ...]:
+ """Serialize the result set, with the stop landing meanwhile."""
+ if stop_at == "serialize":
+ deliver_stop()
+ return ([], [], [], [])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(sql_lab, "_serialize_and_expand_data", side_effect=serialize)
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob, with the stop landing meanwhile."""
+ stored[key] = value
+ if stop_at == "backend_write":
+ deliver_stop()
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="select-stop",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ payload = sql_lab.get_sql_results(
+ 123,
+ query.sql,
+ return_results=return_results,
+ store_results=store_results,
+ )
+
+ assert payload == {"query_id": 123, "status": QueryStatus.STOPPED}
+ session.refresh(query)
+ assert query.status == QueryStatus.STOPPED
+ assert query.end_time == 1234.5
+ assert query.results_key is None
+ if store_results:
+ key = backend.set.call_args.args[0]
+ backend.delete.assert_called_once_with(key)
+ assert key not in stored
+ else:
+ backend.set.assert_not_called()
+ backend.delete.assert_not_called()
+ engine.dispose()
+
+
+@pytest.mark.parametrize("delete_fails", [False, True], ids=["deleted", "delete_error"])
+@pytest.mark.parametrize("return_results", [False, True])
+@pytest.mark.parametrize("stopped", [False, True], ids=["no_stop", "stop"])
+@pytest.mark.parametrize("failure", ["write_timing", "final_flush"])
+def test_get_sql_results_error_after_write_deletes_results_when_stopped( # noqa: C901
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ caplog: pytest.LogCaptureFixture,
+ failure: str,
+ stopped: bool,
+ return_results: bool,
+ delete_fails: bool,
+) -> None:
+ """
+ Test that results stored before a failure are deleted when the error
+ handler finds the query stopped, that a failed delete still ends STOPPED,
+ and that without a stop the failed row keeps the key, even when the failure
+ rolls back the session.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'error-stored.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = (
+ RuntimeError("Backend unavailable") if delete_fails else delete
+ )
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ failed = False
+
+ def fail() -> None:
+ """Raise once the results are stored, with the stop landing first."""
+ nonlocal failed
+ assert stored
+ failed = True
+ if stopped:
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stop_session.commit()
+ raise RuntimeError("Metrics or metadata database unavailable")
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the write timing."""
+ if failure == "write_timing" and key == "sqllab.query.results_backend_write":
+ if not failed:
+ fail()
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def fail_flush( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Fail flushing the result metadata, which rolls back the session."""
+ update = statement.lstrip().startswith("UPDATE")
+ if failure == "final_flush" and update and stored and not failed:
+ fail()
+
+ # get_query retries after rolling back the session the failed flush broke.
+ mocker.patch("backoff._sync.time.sleep")
+ stats_logger = MagicMock()
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(app.config, {"STATS_LOGGER": stats_logger, "SQL_MAX_ROW": 0})
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="error-stored",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ with caplog.at_level("WARNING", logger=sql_lab.logger.name):
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=return_results, store_results=True
+ )
+
+ assert failed
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ # The worker's session ends without committing pending metadata.
+ session.rollback()
+ session.refresh(query)
+ assert payload is not None
+ if not stopped:
+ # Not a stop: the failure is reported and the row keeps its key.
+ assert payload["status"] == QueryStatus.FAILED
+ assert query.status == QueryStatus.FAILED
+ assert query.results_key == key
+ assert key in stored
+ backend.delete.assert_not_called()
+ else:
+ assert payload["status"] == QueryStatus.STOPPED
+ assert query.status == QueryStatus.STOPPED
+ assert query.results_key is None
+ backend.delete.assert_called_once_with(key)
+ assert (key in stored) is delete_fails
+ assert (
+ "Failed to delete unpublished results" in caplog.text
+ ) is delete_fails
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ ("reread", "delete_fails"),
+ [("available", False), ("available", True), ("unavailable", False)],
+ ids=["reread-deleted", "reread-delete_error", "no_reread"],
+)
+def test_get_sql_results_unrecorded_error_after_write_deletes_unreferenced_results(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ caplog: pytest.LogCaptureFixture,
+ reread: str,
+ delete_fails: bool,
+) -> None:
+ """
+ Test that results stored before a failure the task cannot record are
+ deleted once a fresh read shows the row does not reference them, and kept
+ until they expire when the metadata database cannot be read at all.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'unrecorded.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+ mocker.patch("backoff._sync.time.sleep")
+ mocker.patch.dict(app.config, {"SQL_MAX_ROW": 0})
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = (
+ RuntimeError("Backend unavailable") if delete_fails else delete
+ )
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def metadata_unavailable( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Fail every metadata statement once the results are stored."""
+ cleanup_read = statement.startswith('SELECT "query".results_key ')
+ if stored and not (reread == "available" and cleanup_read):
+ raise RuntimeError("Metadata database unavailable")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="unrecorded",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ with caplog.at_level("WARNING", logger=sql_lab.logger.name):
+ with pytest.raises(sql_lab.SqlLabException):
+ sql_lab.get_sql_results(
+ 123, query.sql, return_results=False, store_results=True
+ )
+
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ if reread == "available":
+ backend.delete.assert_called_once_with(key)
+ assert (key in stored) is delete_fails
+ assert (
+ "Failed to delete unpublished results" in caplog.text
+ ) is delete_fails
+ else:
+ backend.delete.assert_not_called()
+ assert key in stored
+ assert "Kept results that may be referenced" in caplog.text
+ engine.dispose()
+
+
+@pytest.mark.parametrize(
+ ("applied", "return_results"),
+ [("failed_row", False), ("failed_row", True), ("published_row", True)],
+)
+def test_get_sql_results_commit_that_applies_and_raises_keeps_results( # noqa: C901
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ applied: str,
+ return_results: bool,
+) -> None:
+ """
+ Test that a metadata commit that reaches the database and then raises (its
+ acknowledgement is lost) does not delete the results the committed row
+ references: the failed row of the error handler, or the published row
+ when the metadata database then stays unavailable.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'applied-commit.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+ mocker.patch("backoff._sync.time.sleep")
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ stats_logger = MagicMock()
+ if applied == "failed_row":
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the write timing, once the results are stored."""
+ if key == "sqllab.query.results_backend_write":
+ raise RuntimeError("Metrics unavailable")
+
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(app.config, {"STATS_LOGGER": stats_logger, "SQL_MAX_ROW": 0})
+
+ target = QueryStatus.FAILED if applied == "failed_row" else QueryStatus.SUCCESS
+ flushed = False
+ applied_commit = False
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def metadata_unavailable( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Keep the metadata database down once the published row commits."""
+ if applied == "published_row" and applied_commit:
+ raise RuntimeError("Metadata database unavailable")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="applied-commit",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ @event.listens_for(session, "after_flush")
+ def flushed_target(flushing_session: Any, context: Any) -> None:
+ nonlocal flushed
+ flushed = query.status == target and query.results_key is not None
+
+ do_commit = engine.dialect.do_commit
+
+ def lose_acknowledgement(dbapi_connection: Any) -> None:
+ """Raise after the database committed the row with the key."""
+ nonlocal applied_commit
+ do_commit(dbapi_connection)
+ if flushed and not applied_commit:
+ applied_commit = True
+ raise RuntimeError("Commit acknowledgement lost")
+
+ mocker.patch.object(
+ engine.dialect, "do_commit", side_effect=lose_acknowledgement
+ )
+
+ with pytest.raises((RuntimeError, sql_lab.SqlLabException)):
+ sql_lab.get_sql_results(
+ 123, query.sql, return_results=return_results, store_results=True
+ )
+
+ assert applied_commit
+ event.remove(engine, "before_cursor_execute", metadata_unavailable)
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == target
+ assert persisted.results_key == key
+ backend.delete.assert_not_called()
+ assert key in stored
+ engine.dispose()
+
+
+@pytest.mark.parametrize("return_results", [False, True])
+def test_get_sql_results_stop_after_an_applied_publish_keeps_results(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ return_results: bool,
+) -> None:
+ """
+ Test that when the publish commit reaches the database and then raises, and
+ a stop lands while the error is handled, the stop wins without deleting the
+ results the committed row references.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'applied-publish.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+ mocker.patch("backoff._sync.time.sleep")
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = lambda key: stored.pop(key, None) is not None
+ mocker.patch.object(sql_lab, "results_backend", backend)
+ mocker.patch.dict(app.config, {"STATS_LOGGER": MagicMock(), "SQL_MAX_ROW": 0})
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="applied-publish",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ published = False
+
+ @event.listens_for(session, "after_flush")
+ def published_flush(flushing_session: Any, context: Any) -> None:
+ nonlocal published
+ published = (
+ query.status == QueryStatus.SUCCESS and query.results_key is not None
+ )
+
+ do_commit = engine.dialect.do_commit
+ applied_commit = False
+
+ def lose_acknowledgement(dbapi_connection: Any) -> None:
+ """Raise after the published row commits and a stop lands."""
+ nonlocal applied_commit
+ do_commit(dbapi_connection)
+ if published and not applied_commit:
+ applied_commit = True
+ with Session(engine) as stop_session:
+ stopped_query = stop_session.get(Query, 123)
+ assert stopped_query is not None
+ stopped_query.status = QueryStatus.STOPPED
+ stop_session.commit()
+ raise RuntimeError("Commit acknowledgement lost")
+
+ mocker.patch.object(
+ engine.dialect, "do_commit", side_effect=lose_acknowledgement
+ )
+
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=return_results, store_results=True
+ )
+
+ assert applied_commit
+ assert payload is not None
+ assert payload["status"] == QueryStatus.STOPPED
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatus.STOPPED
+ assert persisted.results_key == key
+ backend.delete.assert_not_called()
+ assert key in stored
+ engine.dispose()
+
+
+@pytest.mark.parametrize("return_results", [False, True])
+@pytest.mark.parametrize("failure", ["after_failed_commit", "failed_commit"])
+def test_get_sql_results_deletes_results_only_without_committed_failure( # noqa: C901
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ failure: str,
+ return_results: bool,
+) -> None:
+ """
+ Test that results stored before a failure are kept once the failed row
+ referencing them is committed, even if the metadata database fails right
+ after that commit, and are deleted when the failed row's commit fails.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'failed-commit.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ def delete(key: str) -> bool:
+ """Remove the blob."""
+ return stored.pop(key, None) is not None
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ backend.delete.side_effect = delete
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ def timing(key: str, value: float) -> None:
+ """Fail emitting the write timing, once the results are stored."""
+ if key == "sqllab.query.results_backend_write":
+ raise RuntimeError("Metrics unavailable")
+
+ stats_logger = MagicMock()
+ stats_logger.timing.side_effect = timing
+ mocker.patch.dict(app.config, {"STATS_LOGGER": stats_logger, "SQL_MAX_ROW": 0})
+
+ failed_flushed = False
+ failed_committed = False
+
+ @event.listens_for(engine, "commit")
+ def fail_failed_commit(conn: Any) -> None:
+ """Fail committing the failed row before the database commits it."""
+ if failure == "failed_commit" and failed_flushed:
+ raise RuntimeError("Metadata commit failed")
+
+ @event.listens_for(engine, "before_cursor_execute")
+ def fail_after_failed_commit( # noqa: PLR0913
+ conn: Any,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any,
+ executemany: bool,
+ ) -> None:
+ """Fail every metadata statement once the failed row is committed."""
+ if failed_committed:
+ raise RuntimeError("Metadata database unavailable")
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="failed-commit",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ @event.listens_for(session, "after_flush")
+ def flushed_failure(flushing_session: Any, context: Any) -> None:
+ nonlocal failed_flushed
+ failed_flushed = query.status == QueryStatus.FAILED
+
+ @event.listens_for(session, "after_commit")
+ def committed_failure(committing_session: Any) -> None:
+ nonlocal failed_committed
+ failed_committed = failed_flushed
+
+ if failure == "after_failed_commit":
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=return_results, store_results=True
+ )
+ assert payload is not None
+ assert payload["status"] == QueryStatus.FAILED
+ assert failed_committed
+ else:
+ with pytest.raises(RuntimeError, match="Metadata commit failed"):
+ sql_lab.get_sql_results(
+ 123, query.sql, return_results=return_results, store_results=True
+ )
+ assert not failed_committed
+ session.rollback()
+
+ event.remove(engine, "before_cursor_execute", fail_after_failed_commit)
+ backend.set.assert_called_once()
+ key = backend.set.call_args.args[0]
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ if failure == "after_failed_commit":
+ # The committed failed row references the results; keep them.
+ assert persisted.status == QueryStatus.FAILED
+ assert persisted.results_key == key
+ assert key in stored
+ backend.delete.assert_not_called()
+ else:
+ # Nothing committed references the results; delete them.
+ assert persisted.status != QueryStatus.FAILED
+ assert persisted.results_key is None
+ assert key not in stored
+ backend.delete.assert_called_once_with(key)
+ engine.dispose()
+
+
+def test_get_sql_results_keeps_published_results_after_final_commit(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+) -> None:
+ """
+ Test that a failure after the final metadata commit does not delete results
+ whose key the committed query row already references.
+ """
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'published.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ # The stored payload fits; the returned one exceeds SQLLAB_PAYLOAD_MAX_MB,
+ # so its check fails after the final metadata commit.
+ mocker.patch.object(
+ sql_lab,
+ "_serialize_payload",
+ side_effect=[b"data", b"x" * (2 * 1024 * 1024)],
+ )
+ mocker.patch.dict(
+ app.config,
+ {"STATS_LOGGER": MagicMock(), "SQL_MAX_ROW": 0, "SQLLAB_PAYLOAD_MAX_MB": 1},
+ )
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="published",
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=True, store_results=True
+ )
+
+ key = backend.set.call_args.args[0]
+ session.rollback()
+ session.refresh(query)
+ assert payload is not None
+ assert payload["status"] == QueryStatus.FAILED
+ assert query.status == QueryStatus.FAILED
+ assert query.results_key == key
+ assert key in stored
+ backend.delete.assert_not_called()
+ engine.dispose()
+
+
def _capture_execution_context(mocker: MockerFixture) -> dict[str, Any]:
"""
Stub out ``execute_sql_statements`` and record the context it runs under.
@@ -999,3 +2326,390 @@ def test_get_predicates_for_table_excludes_self(mocker: MockerFixture) -> None:
filter_call = db.session.query().filter.call_args
and_clause = filter_call.args[0]
assert len(and_clause.clauses) == 5
+
+
+@pytest.mark.parametrize("return_results", [False, True], ids=["async", "sync"])
+@pytest.mark.parametrize("stop_timing", ["cursor", "error_handler"])
+def test_get_sql_results_preserves_stop_during_postgres_cancellation(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ return_results: bool,
+ stop_timing: str,
+) -> None:
+ """A cancellation failure cannot overwrite a Stop from another session."""
+ from unittest.mock import PropertyMock
+
+ from psycopg2.errors import QueryCanceled
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'cancel-error.db'}")
+ Query.__table__.create(engine)
+ database = MagicMock()
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ database.get_default_schema_for_query.return_value = "public"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.dict(
+ app.config, {"STATS_LOGGER": MagicMock(), "DISALLOWED_SQL_FUNCTIONS": {}}
+ )
+
+ def stop() -> None:
+ """Commit the Stop independently of the worker's stale identity map."""
+ with Session(engine) as stop_session:
+ stopped = stop_session.get(Query, 123)
+ assert stopped is not None
+ stopped.status = QueryStatus.STOPPED
+ stopped.end_time = 1234.5
+ stop_session.commit()
+
+ cursor_cancelled = False
+
+ def cancelled_cursor(*args: Any, **kwargs: Any) -> None:
+ """Simulate Postgres raising after the running statement is cancelled."""
+ nonlocal cursor_cancelled
+ cursor_cancelled = True
+ raise QueryCanceled("canceling statement due to user request")
+
+ def extract_errors(*args: Any, **kwargs: Any) -> list[Any]:
+ """Let Stop commit after the error handler's preliminary status read."""
+ if stop_timing == "error_handler":
+ stop()
+ return []
+
+ mocker.patch.object(sql_lab, "warm_and_release_connection")
+
+ # Use the real legacy execute_query and error handler, not a task stub.
+ mocker.patch.object(PostgresEngineSpec, "get_cancel_query_id", return_value=42)
+ execute_cursor = mocker.patch.object(
+ PostgresEngineSpec, "execute_with_cursor", side_effect=cancelled_cursor
+ )
+ mocker.patch.object(sql_lab, "_extract_errors", side_effect=extract_errors)
+ with Session(engine) as worker_session:
+ query = Query(
+ id=123,
+ client_id="cancel-race",
+ database_id=1,
+ sql="SELECT pg_sleep(30)",
+ status=QueryStatus.PENDING,
+ )
+ worker_session.add(query)
+ worker_session.commit()
+ mocker.patch.object(sql_lab.db, "session", worker_session)
+ refresh = worker_session.refresh
+
+ def refresh_then_stop(*args: Any, **kwargs: Any) -> None:
+ """Commit after the cancelled cursor's handler reads its status."""
+ refresh(*args, **kwargs)
+ if (
+ stop_timing == "cursor"
+ and cursor_cancelled
+ and kwargs.get("attribute_names") == ["status"]
+ ):
+ stop()
+
+ mocker.patch.object(worker_session, "refresh", side_effect=refresh_then_stop)
+ payload = sql_lab.get_sql_results(123, query.sql, return_results=return_results)
+ execute_cursor.assert_called_once()
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatus.STOPPED
+ assert persisted.end_time == 1234.5
+ assert persisted.error_message is None
+ assert payload is not None
+ assert payload["status"] == QueryStatus.STOPPED
+ assert "error" not in payload
+ engine.dispose()
+
+
+@pytest.mark.parametrize("mutation", [False, True], ids=["select", "mutation"])
+@pytest.mark.parametrize("terminal_write", ["success", "backend_failure"])
+def test_legacy_terminal_write_preserves_concurrent_stop(
+ mocker: MockerFixture,
+ request: pytest.FixtureRequest,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+ terminal_write: str,
+ mutation: bool,
+) -> None:
+ """Neither result publication nor a cache failure may overwrite a Stop."""
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine, event
+
+ from superset import sql_lab
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'terminal-stop.db'}")
+ Query.__table__.create(engine)
+ database = MagicMock()
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ database.allow_dml = True
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "public"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+ mocker.patch.object(sql_lab, "_serialize_payload", return_value=b"data")
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ mocker.patch.dict(app.config, {"STATS_LOGGER": MagicMock(), "SQL_MAX_ROW": 0})
+
+ def stop() -> None:
+ """Publish Stop through a separate metadata transaction."""
+ with Session(engine) as stop_session:
+ stopped = stop_session.get(Query, 123)
+ assert stopped is not None
+ stopped.status = QueryStatus.STOPPED
+ stopped.end_time = 1234.5
+ stop_session.commit()
+
+ def store(*args: Any) -> bool:
+ """Return a failure before racing Stop with its terminal write."""
+ return terminal_write != "backend_failure"
+
+ stop_attempted = False
+
+ def stop_before_failure(
+ target: Query, value: str, old_value: Any, initiator: Any
+ ) -> None:
+ """Commit after any preliminary refresh, before FAILED is flushed."""
+ nonlocal stop_attempted
+ if (
+ terminal_write == "backend_failure"
+ and value == QueryStatus.FAILED
+ and not stop_attempted
+ ):
+ stop_attempted = True
+ stop()
+
+ event.listen(Query.status, "set", stop_before_failure)
+ request.addfinalizer(lambda: event.remove(Query.status, "set", stop_before_failure))
+
+ mocker.patch.object(PostgresEngineSpec, "get_cancel_query_id", return_value=42)
+ backend = MagicMock()
+ backend.set.side_effect = store
+ mocker.patch.object(sql_lab, "results_backend", backend)
+ with Session(engine) as worker_session:
+ query = Query(
+ id=123,
+ client_id="final-stop",
+ database_id=1,
+ sql="INSERT INTO users VALUES (1)" if mutation else "SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ worker_session.add(query)
+ worker_session.commit()
+ mocker.patch.object(sql_lab.db, "session", worker_session)
+
+ def execute(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit execution metadata as the real statement executor does."""
+ worker_session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute)
+ refresh = worker_session.refresh
+
+ def refresh_then_stop(*args: Any, **kwargs: Any) -> None:
+ """Stop between the final status read and terminal publication."""
+ refresh(*args, **kwargs)
+ if (
+ terminal_write == "success"
+ and kwargs.get("attribute_names") == ["status"]
+ and query.results_key
+ ):
+ stop()
+
+ mocker.patch.object(worker_session, "refresh", side_effect=refresh_then_stop)
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=False, store_results=True
+ )
+ backend.set.assert_called_once()
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatus.STOPPED
+ assert persisted.end_time == 1234.5
+ assert persisted.results_key is None
+ assert persisted.error_message is None
+ assert payload == {"query_id": 123, "status": QueryStatus.STOPPED}
+ if terminal_write == "success":
+ backend.delete.assert_called_once_with(backend.set.call_args.args[0])
+ engine.dispose()
+
+
+def test_legacy_startup_preserves_stop_before_running_write(
+ mocker: MockerFixture,
+ app_context: None,
+ tmp_path: Any,
+) -> None:
+ """A Stop between the startup read and RUNNING write prevents dispatch."""
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'startup-stop.db'}")
+ Query.__table__.create(engine)
+ database = MagicMock()
+ database.db_engine_spec = PostgresEngineSpec
+ database.allow_run_async = False
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+
+ def stop() -> None:
+ """Commit after the worker's initial status check."""
+ with Session(engine) as stop_session:
+ stopped = stop_session.get(Query, 123)
+ assert stopped is not None
+ stopped.status = QueryStatus.STOPPED
+ stopped.end_time = 1234.5
+ stop_session.commit()
+
+ mocker.patch.object(PostgresEngineSpec, "patch", side_effect=stop)
+ parse = mocker.patch.object(
+ sql_lab,
+ "SQLScript",
+ side_effect=AssertionError("Stopped query reached parsing"),
+ )
+ with Session(engine) as worker_session:
+ worker_session.add(
+ Query(
+ id=123, client_id="startup", database_id=1, status=QueryStatus.PENDING
+ )
+ )
+ worker_session.commit()
+ mocker.patch.object(sql_lab.db, "session", worker_session)
+ payload = sql_lab.execute_sql_statements(
+ 123, "SELECT 1", True, False, None, False, None
+ )
+ assert payload == {"query_id": 123, "status": QueryStatus.STOPPED}
+ parse.assert_not_called()
+ with Session(engine) as observer:
+ persisted = observer.get(Query, 123)
+ assert persisted is not None
+ assert persisted.status == QueryStatus.STOPPED
+ assert persisted.end_time == 1234.5
+ engine.dispose()
+
+
+def test_sync_payload_limit_stop_response_matches_recorded_outcome(
+ mocker: MockerFixture,
+ app: SupersetApp,
+ app_context: None,
+ tmp_path: Any,
+) -> None:
+ """Stop during inline size validation cannot observe a premature SUCCESS."""
+ from unittest.mock import PropertyMock
+
+ from sqlalchemy import create_engine
+
+ from superset import sql_lab
+ from superset.db_engine_specs.sqlite import SqliteEngineSpec
+ from superset.models.sql_lab import Query
+
+ engine = create_engine(f"sqlite:///{tmp_path / 'published.db'}")
+ Query.__table__.create(engine)
+
+ database = MagicMock()
+ database.db_engine_spec = SqliteEngineSpec
+ database.allow_run_async = False
+ database.cache_timeout = 100
+ database.get_default_schema_for_query.return_value = "main"
+ database.mutate_sql_based_on_config.side_effect = lambda sql, is_split=False: sql
+ mocker.patch.object(Query, "database", PropertyMock(return_value=database))
+ mocker.patch.object(Query, "to_dict", return_value={"id": 123})
+ mocker.patch.object(sql_lab.security_manager, "find_user", return_value=None)
+
+ def execute_query(*args: Any, **kwargs: Any) -> MagicMock:
+ """Commit the pre-execution metadata as execute_query does."""
+ sql_lab.db.session.commit()
+ return MagicMock(size=1, columns=[])
+
+ mocker.patch.object(sql_lab, "execute_query", side_effect=execute_query)
+ mocker.patch.object(
+ sql_lab, "_serialize_and_expand_data", return_value=([], [], [], [])
+ )
+ # The stored payload fits; the returned one exceeds SQLLAB_PAYLOAD_MAX_MB,
+ # so the inline size check must run before publishing SUCCESS.
+ stop_responses: list[int] = []
+ serializations = 0
+
+ def serialize_with_stop(*args: Any) -> bytes:
+ """Stop at the inline payload check, before any oversized-result error."""
+ import inspect
+
+ from superset.queries.api import QueryRestApi
+
+ nonlocal serializations
+ serializations += 1
+ if serializations == 1:
+ return b"data"
+ with app.test_request_context(json={"client_id": "published"}):
+ response = inspect.unwrap(QueryRestApi.stop_query)(QueryRestApi())
+ stop_responses.append(response.status_code)
+ return b"x" * (2 * 1024 * 1024)
+
+ mocker.patch("superset.daos.query.get_user_id", return_value=7)
+ mocker.patch.object(sql_lab, "cancel_query", return_value=False)
+ mocker.patch.object(
+ sql_lab,
+ "_serialize_payload",
+ side_effect=serialize_with_stop,
+ )
+ mocker.patch.dict(
+ app.config,
+ {"STATS_LOGGER": MagicMock(), "SQL_MAX_ROW": 0, "SQLLAB_PAYLOAD_MAX_MB": 1},
+ )
+
+ stored: dict[str, bytes] = {}
+
+ def store(key: str, value: bytes, timeout: int) -> bool:
+ """Write the blob successfully."""
+ stored[key] = value
+ return True
+
+ backend = MagicMock()
+ backend.set.side_effect = store
+ mocker.patch.object(sql_lab, "results_backend", backend)
+
+ with Session(engine) as session:
+ query = Query(
+ id=123,
+ client_id="published",
+ user_id=7,
+ database_id=1,
+ sql="SELECT 1",
+ executed_sql="SELECT 1",
+ status=QueryStatus.PENDING,
+ )
+ session.add(query)
+ session.commit()
+ mocker.patch.object(sql_lab.db, "session", session)
+
+ payload = sql_lab.get_sql_results(
+ 123, query.sql, return_results=True, store_results=True
+ )
+
+ key = backend.set.call_args.args[0]
+ session.rollback()
+ session.refresh(query)
+ assert payload is not None
+ assert stop_responses == [200]
+ assert payload["status"] == QueryStatus.STOPPED
+ assert query.status == QueryStatus.STOPPED
+ assert query.results_key is None
+ backend.delete.assert_called_once_with(key)
+ engine.dispose()
diff --git a/tests/unit_tests/test_sqllab_error_stacktrace.py b/tests/unit_tests/test_sqllab_error_stacktrace.py
index 9f93f0a2353..83d835b30d0 100644
--- a/tests/unit_tests/test_sqllab_error_stacktrace.py
+++ b/tests/unit_tests/test_sqllab_error_stacktrace.py
@@ -53,6 +53,7 @@ def _make_query_mock() -> MagicMock:
"""Return a minimal Query mock that satisfies handle_query_error."""
query = MagicMock()
query.status = "running"
+ query.extra = {}
query.end_time = None
query.tmp_table_name = "tmp"
query.error_message = None
@@ -78,6 +79,7 @@ def test_legacy_handle_query_error_payload_includes_stacktrace_when_enabled() ->
config = {"TROUBLESHOOTING_LINK": None, "SHOW_STACKTRACE": True}
with (
patch("superset.sql_lab.db") as mock_db,
+ patch("superset.sql.execution.utils.db", mock_db),
patch("superset.sql_lab.app") as mock_app,
):
mock_app.config = config
@@ -117,6 +119,7 @@ def test_legacy_handle_query_error_payload_omits_stacktrace_when_disabled() -> N
config = {"TROUBLESHOOTING_LINK": None, "SHOW_STACKTRACE": False}
with (
patch("superset.sql_lab.db") as mock_db,
+ patch("superset.sql.execution.utils.db", mock_db),
patch("superset.sql_lab.app") as mock_app,
):
mock_app.config = config
@@ -197,7 +200,11 @@ def test_new_handle_query_error_payload_includes_stacktrace_when_enabled() -> No
query = _make_query_mock()
config = {"TROUBLESHOOTING_LINK": None, "SHOW_STACKTRACE": True}
- with patch.object(ct, "db") as mock_db, patch.object(ct, "app") as mock_app:
+ with (
+ patch.object(ct, "db") as mock_db,
+ patch("superset.sql.execution.utils.db", mock_db),
+ patch.object(ct, "app") as mock_app,
+ ):
mock_app.config = config
mock_db.session = MagicMock()
@@ -206,6 +213,7 @@ def test_new_handle_query_error_payload_includes_stacktrace_when_enabled() -> No
except RuntimeError as ex:
payload = ct._handle_query_error(ex, query)
+ assert payload["status"] == "failed"
assert "stacktrace" in payload, (
"_handle_query_error must include 'stacktrace' in the returned payload "
"when SHOW_STACKTRACE=True (regression for #28248). "
@@ -234,7 +242,11 @@ def test_new_handle_query_error_payload_omits_stacktrace_when_disabled() -> None
query = _make_query_mock()
config = {"TROUBLESHOOTING_LINK": None, "SHOW_STACKTRACE": False}
- with patch.object(ct, "db") as mock_db, patch.object(ct, "app") as mock_app:
+ with (
+ patch.object(ct, "db") as mock_db,
+ patch("superset.sql.execution.utils.db", mock_db),
+ patch.object(ct, "app") as mock_app,
+ ):
mock_app.config = config
mock_db.session = MagicMock()
@@ -243,6 +255,7 @@ def test_new_handle_query_error_payload_omits_stacktrace_when_disabled() -> None
except RuntimeError as ex:
payload = ct._handle_query_error(ex, query)
+ assert payload["status"] == "failed"
assert "stacktrace" not in payload, (
"_handle_query_error must NOT include 'stacktrace' when SHOW_STACKTRACE=False "
"to avoid exposing internal details to unprivileged users. "