aminghadersohi commented on code in PR #45283:
URL: https://github.com/apache/superset/pull/45283#discussion_r4278722639
##########
superset/sql/execution/celery_task.py:
##########
@@ -198,102 +218,110 @@ def _finalize_successful_query(
}
)
- query.rows = total_rows
- query.progress = 100
- query.set_extra_json_key("progress", None)
- # Store columns from last statement (for compatibility)
- if execution_results and execution_results[-1][1] is not None:
- query.set_extra_json_key("columns", execution_results[-1][1].columns)
- query.end_time = now_as_float()
-
- payload.update(
- {
- "status": QueryStatus.SUCCESS.value,
- "statements": statements_data,
- "total_execution_time_ms": total_execution_time_ms,
- "query": query.to_dict(),
- }
- )
- payload["query"]["state"] = QueryStatus.SUCCESS.value
+ with db.session.no_autoflush:
+ query.rows = total_rows
+ query.progress = 100
+ query.set_extra_json_key("progress", None)
+ # Store columns from last statement (for compatibility)
+ if execution_results and execution_results[-1][1] is not None:
+ query.set_extra_json_key("columns",
execution_results[-1][1].columns)
+ query.end_time = now_as_float()
+
+ payload.update(
+ {
+ "status": QueryStatus.SUCCESS.value,
+ "statements": statements_data,
+ "total_execution_time_ms": total_execution_time_ms,
+ "query": query.to_dict(),
+ }
+ )
+ payload["query"]["state"] = QueryStatus.SUCCESS.value
def _store_results_in_backend(
query: Query,
payload: dict[str, Any],
database: Any,
+ durable: bool = False,
) -> None:
"""Store query results in the results backend."""
- key = str(uuid.uuid4())
- payload["query"]["resultsKey"] = key
- logger.info(
- "Query %s: Storing results in results backend, key: %s",
- str(query.id),
- key,
- )
- stats_logger = app.config["STATS_LOGGER"]
- with stats_timing("sqllab.query.results_backend_write", stats_logger):
- with stats_timing(
- "sqllab.query.results_backend_write_serialization", stats_logger
- ):
- serialized_payload = _serialize_payload(payload)
-
- # Check payload size limit
- if sql_lab_payload_max_mb :=
app.config.get("SQLLAB_PAYLOAD_MAX_MB"):
- serialized_payload_size = len(serialized_payload)
- max_bytes = sql_lab_payload_max_mb * BYTES_IN_MB
-
- if serialized_payload_size > max_bytes:
- logger.info("Result size exceeds the allowed limit.")
- raise SupersetErrorException(
- SupersetError(
- message=(
- f"Result size "
- f"({serialized_payload_size / BYTES_IN_MB:.2f}
MB) "
- f"exceeds the allowed limit of "
- f"{sql_lab_payload_max_mb} MB."
- ),
-
error_type=SupersetErrorType.RESULT_TOO_LARGE_ERROR,
- level=ErrorLevel.ERROR,
+ with db.session.no_autoflush:
+ key = str(uuid.uuid4())
+ payload["query"]["resultsKey"] = key
+ logger.info(
+ "Query %s: Storing results in results backend, key: %s",
+ str(query.id),
+ key,
+ )
+ stats_logger = app.config["STATS_LOGGER"]
+ with stats_timing("sqllab.query.results_backend_write", stats_logger):
+ with stats_timing(
+ "sqllab.query.results_backend_write_serialization",
stats_logger
+ ):
+ serialized_payload = _serialize_payload(payload)
+
+ # Check payload size limit
+ if sql_lab_payload_max_mb :=
app.config.get("SQLLAB_PAYLOAD_MAX_MB"):
+ serialized_payload_size = len(serialized_payload)
+ max_bytes = sql_lab_payload_max_mb * BYTES_IN_MB
+
+ if serialized_payload_size > max_bytes:
+ logger.info("Result size exceeds the allowed limit.")
+ raise SupersetErrorException(
+ SupersetError(
+ message=(
+ f"Result size "
+ f"({serialized_payload_size /
BYTES_IN_MB:.2f} MB) "
+ f"exceeds the allowed limit of "
+ f"{sql_lab_payload_max_mb} MB."
+ ),
+
error_type=SupersetErrorType.RESULT_TOO_LARGE_ERROR,
+ level=ErrorLevel.ERROR,
+ )
)
- )
-
- cache_timeout = database.cache_timeout
- if cache_timeout is None:
- cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"]
- compressed = zlib_compress(serialized_payload)
- logger.debug("*** serialized payload size: %i",
len(serialized_payload))
- logger.debug("*** compressed payload size: %i", len(compressed))
-
- write_success = results_backend.set(key, compressed, cache_timeout)
- if not write_success:
- logger.error(
- "Query %s: Failed to store results in backend, key: %s",
- str(query.id),
- key,
- )
- stats_logger.incr("sqllab.results_backend.write_failure")
- query.results_key = None
- query.status = QueryStatus.FAILED
- query.error_message = (
- "Failed to store query results in the results backend. "
- "Please try again or contact your administrator."
- )
- db.session.commit() # pylint: disable=consider-using-transaction
- raise SupersetErrorException(
- SupersetError(
- message=__("Failed to store query results. Please try
again."),
- error_type=SupersetErrorType.RESULTS_BACKEND_ERROR,
- level=ErrorLevel.ERROR,
+ cache_timeout = database.cache_timeout
+ if cache_timeout is None:
+ cache_timeout = app.config["CACHE_DEFAULT_TIMEOUT"]
+
+ compressed = zlib_compress(serialized_payload)
+ logger.debug("*** serialized payload size: %i",
len(serialized_payload))
+ logger.debug("*** compressed payload size: %i", len(compressed))
+
+ write_success = results_backend.set(key, compressed, cache_timeout)
+ if not durable and _make_check_stopped_fn(query,
preserve_pending=True)():
Review Comment:
Follow-up: an independent check showed my last reply was wrong about three
siblings I listed as not affected (superset/sql/execution/celery_task.py:106,
superset/sql/execution/celery_task.py:137 and superset/sql_lab.py:151). A
publish commit can reach the database and still raise. If a stop then landed
while that error was handled, those stop branches deleted the blob directly,
and the committed STOPPED row kept a `results_key` pointing at deleted results
(Celery, and the legacy path in both `return_results` modes).
Fixed in e53b6f823fef8785bef50088a754b659d417cc7e and
bb19cf7abf62d8519c2766648e05729ef1834da3. All three stop branches
(superset/sql/execution/celery_task.py:106,
superset/sql/execution/celery_task.py:137, superset/sql_lab.py:154) now use the
same rule as the unrecorded-failure cleanup: `_discard_unreferenced_results`
rolls back, reads the row's `results_key` in a new transaction, and deletes the
blob only if the row does not reference it. If that read fails, the blob is
kept until `cache_timeout`. The first commit compared against the row loaded in
the worker's own transaction; the legacy `get_query` autoflushes, so that read
saw the worker's uncommitted key. bb19cf7abf62d8519c2766648e05729ef1834da3
switches to the fresh read.
Regression tests. Each fails at 0eab948613546ed4a7a678c7b19c5fcdc3220ab7
with "Expected 'delete' to not have been called. Called 1 times." and passes at
bb19cf7abf62d8519c2766648e05729ef1834da3. Each asserts the payload is STOPPED,
a second session reads a STOPPED row whose `results_key` is the stored key, and
the blob is still stored:
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_stop_after_an_applied_publish_keeps_referenced_results[before_stop_check]`:
the stop is seen by the first check (line 106).
- `[before_failed_commit]` of the same test: the stop lands after that check
and is seen by the guarded commit (line 137).
-
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_stop_after_an_applied_publish_keeps_results[False]`
and `[True]`: the legacy handler, in both `return_results` modes.
- Controls that pass on both heads:
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_error_after_write_deletes_results_when_stopped`
and
`tests/unit_tests/sql/execution/test_celery_task.py::test_error_after_results_write_deletes_stored_results_when_stopped`
(a stop with no committed reference still deletes the blob, also when the
delete fails).
Siblings checked:
- superset/sql/execution/celery_task.py:106 — fixed in
bb19cf7abf62d8519c2766648e05729ef1834da3
- superset/sql/execution/celery_task.py:137 — fixed in
bb19cf7abf62d8519c2766648e05729ef1834da3
- superset/sql_lab.py:154 — fixed in bb19cf7abf62d8519c2766648e05729ef1834da3
- superset/sql/execution/celery_task.py:451 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 (already used this rule)
- superset/sql_lab.py:370 — fixed in
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 (already used this rule)
- superset/sql/execution/celery_task.py:315 — not affected: the stop is seen
right after the write, before any commit carries the key
- superset/sql/execution/celery_task.py:762 — not affected: the stop is seen
before the publish commit
- superset/sql/execution/celery_task.py:768 — not affected: the guarded
UPDATE matched no row, so the publish committed nothing
- superset/sql_lab.py:981 — not affected: the key is still pending (no flush
since it was set), the reload drops it, and the STOPPED commit does not carry it
Remaining:
- If the metadata database cannot be read after the failure or the stop, the
blob is kept until `cache_timeout`, even when no row references it.
- If the worker dies between the write and the publish, or a delete fails,
the blob also stays until `cache_timeout`, as before.
##########
superset/sql/execution/celery_task.py:
##########
@@ -481,25 +570,60 @@ def _execute_sql_statements(
)
) from ex
- # Check if stopped
- if not execution_results:
- payload.update({"status": QueryStatus.STOPPED.value})
- return payload
+ # Cancelling can clear cursor.description, leaving a nonempty result
+ # list that looks like successful DML. Check intent, not just results.
+ stopped = not execution_results or check_stopped()
# Commit for mutations
- if parsed_script.has_mutation() or query.select_as_cta:
+ committed = False
+ if not stopped and (parsed_script.has_mutation() or
query.select_as_cta):
Review Comment:
Follow-up: an independent check of my last reply found that native SQLite
was not covered, and that the remaining cases were not all listed. This reply
corrects both.
SQLite: native SQLite in its default transaction mode commits DDL, including
a CTAS, as soon as it finishes, so a CTAS that finished before a stop ended
STOPPED although the table was created. Fixed in
6280bdc4d625150d1c4ae3315c5cb0b3b6928317 (head
bb19cf7abf62d8519c2766648e05729ef1834da3). `SqliteEngineSpec.is_autocommit`
(superset/db_engine_specs/sqlite.py:132) reads the driver's `in_transaction`
once the statements have run: no open transaction means nothing is left for
`conn.commit()` to commit or discard. DML in the default mode keeps its
transaction, so a stop before the commit still discards it and ends STOPPED.
Connections that do not expose `in_transaction` fall back to the base check.
Regression tests (each fails at 0eab948613546ed4a7a678c7b19c5fcdc3220ab7 and
passes at bb19cf7abf62d8519c2766648e05729ef1834da3):
-
`tests/unit_tests/db_engine_specs/test_sqlite.py::test_is_autocommit_reports_uncommitted_work[ctas]`,
`[ddl]` and `[dml_isolation_none]`, and
`test_is_autocommit_reads_pooled_connections`, on real sqlite3 connections.
`[dml]` passes on both heads.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_sqlite_ctas_finished_before_a_stop_is_committed[ctas]`
and
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_sqlite_ctas_finished_before_a_stop_is_committed[ctas]`:
a real SQLite target runs the CTAS, a second session commits the stop, and the
query now ends SUCCESS with the table present. At
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 it ended STOPPED. The `[insert]` cases
pass on both heads: the INSERT is discarded and the query ends STOPPED.
Siblings checked:
- superset/db_engine_specs/sqlite.py:132 — fixed in
6280bdc4d625150d1c4ae3315c5cb0b3b6928317
- superset/sql/execution/celery_task.py:704 — not affected: calls
`db_engine_spec.is_autocommit(conn)`, which covers SQLite through the spec
- superset/sql_lab.py:767 — not affected: same call as the Celery path
- superset/db_engine_specs/shillelagh.py:35 — not affected: the Shillelagh
connection (also used by the GSheets and Superset specs) has no
`in_transaction`, so it falls back to the base check as before
- superset/db_engine_specs/d1.py:48 — not affected: the
sqlalchemy-cloudflare-d1 connection has no `in_transaction`, so it falls back
to the base check as before
Remaining (not fixed in this PR; tracked as a follow-up):
- Drivers that report no transaction mode, on engines with no
`autocommit_by_default` setting (for example BigQuery and Snowflake, and the
Shillelagh, GSheets and D1 specs above), can still end STOPPED after durable
writes when the stop is recorded after the statement finished but before the
target-commit check.
- Interrupted autocommit scripts: on a connection that commits each
statement, a stop between statements keeps the statements that already
finished, and the query ends STOPPED without per-statement status.
- Poll-to-cancel: if the engine finishes a statement between its last status
poll and the cancel call, the query ends STOPPED with the writes durable.
- Legacy per-block check to reload window: a stop committed after the
per-block check but before `execute_query` reloads the row still runs the
statement on engines that do not poll for stops, and the query ends STOPPED
with the writes durable.
##########
superset/sql_lab.py:
##########
@@ -628,8 +651,12 @@ def execute_sql_statements( # noqa: C901
return payload
# Commit the connection so CTA queries will create the table and any
DML.
+ committed = False
if parsed_script.has_mutation() or query.select_as_cta:
Review Comment:
Follow-up: an independent check of my last reply found that a finished CTAS
on native SQLite could still end STOPPED, and that my Remaining list was
incomplete. This reply corrects both.
SQLite: native SQLite in its default transaction mode commits DDL, including
a CTAS, as soon as it finishes, so a CTAS that finished before a stop ended
STOPPED although the table was created. Fixed in
6280bdc4d625150d1c4ae3315c5cb0b3b6928317 (head
bb19cf7abf62d8519c2766648e05729ef1834da3). `SqliteEngineSpec.is_autocommit`
(superset/db_engine_specs/sqlite.py:132) reads the driver's `in_transaction`
once the statements have run: no open transaction means nothing is left for
`conn.commit()` to commit or discard. DML in the default mode keeps its
transaction, so a stop before the commit still discards it and ends STOPPED.
Connections that do not expose `in_transaction` fall back to the base check.
Regression tests (each fails at 0eab948613546ed4a7a678c7b19c5fcdc3220ab7 and
passes at bb19cf7abf62d8519c2766648e05729ef1834da3):
-
`tests/unit_tests/db_engine_specs/test_sqlite.py::test_is_autocommit_reports_uncommitted_work[ctas]`,
`[ddl]` and `[dml_isolation_none]`, and
`test_is_autocommit_reads_pooled_connections`, on real sqlite3 connections.
`[dml]` passes on both heads.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_sqlite_ctas_finished_before_a_stop_is_committed[ctas]`
and
`tests/unit_tests/sql_lab_test.py::test_get_sql_results_sqlite_ctas_finished_before_a_stop_is_committed[ctas]`:
a real SQLite target runs the CTAS, a second session commits the stop, and the
query now ends SUCCESS with the table present. At
0eab948613546ed4a7a678c7b19c5fcdc3220ab7 it ended STOPPED. The `[insert]` cases
pass on both heads: the INSERT is discarded and the query ends STOPPED.
Siblings checked:
- superset/db_engine_specs/sqlite.py:132 — fixed in
6280bdc4d625150d1c4ae3315c5cb0b3b6928317
- superset/sql/execution/celery_task.py:704 — not affected: calls
`db_engine_spec.is_autocommit(conn)`, which covers SQLite through the spec
- superset/sql_lab.py:767 — not affected: same call as the Celery path
- superset/db_engine_specs/shillelagh.py:35 — not affected: the Shillelagh
connection (also used by the GSheets and Superset specs) has no
`in_transaction`, so it falls back to the base check as before
- superset/db_engine_specs/d1.py:48 — not affected: the
sqlalchemy-cloudflare-d1 connection has no `in_transaction`, so it falls back
to the base check as before
Remaining (not fixed in this PR; tracked as a follow-up):
- Drivers that report no transaction mode, on engines with no
`autocommit_by_default` setting (for example BigQuery and Snowflake, and the
Shillelagh, GSheets and D1 specs above), can still end STOPPED after durable
writes when the stop is recorded after the statement finished but before the
target-commit check.
- Interrupted autocommit scripts: on a connection that commits each
statement, a stop between statements keeps the statements that already
finished, and the query ends STOPPED without per-statement status.
- Poll-to-cancel: if the engine finishes a statement between its last status
poll and the cancel call, the query ends STOPPED with the writes durable.
- Legacy per-block check to reload window: a stop committed after the
per-block check but before `execute_query` reloads the row still runs the
statement on engines that do not poll for stops, and the query ends STOPPED
with the writes durable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
