aminghadersohi opened a new pull request, #45348:
URL: https://github.com/apache/superset/pull/45348
### SUMMARY
`AsyncQueryHandle.get_result()` (returned by `Database.execute_async()` /
`SQLExecutor.execute_async()`) could not read back any result written by the
async Celery task, so every successful async query was reported as `FAILED`
with no statements:
- **MessagePack payloads** (`RESULTS_BACKEND_USE_MSGPACK = True`, the
default): the reader called `msgpack.loads(zlib_decompress(blob))`.
`zlib_decompress()` defaults to `decode=True`, so it tried to UTF-8 decode the
binary MessagePack map and failed with `'utf-8' codec can't decode byte 0x85`.
- **Arrow statement data**: past that first error, each SELECT statement's
`data` is Arrow IPC bytes, which the reader handed straight to
`pd.DataFrame(...)`.
- **JSON payloads** (`RESULTS_BACKEND_USE_MSGPACK = False`): the decoded
JSON text went to `msgpack.loads()`, which failed with `a bytes-like object is
required, not 'str'`.
- A zero-row SELECT returned `data=None` rather than an empty DataFrame with
its columns, unlike the synchronous path.
The existing unit tests mocked `zlib_decompress` and fed plain records, so
neither side of the real writer/reader contract was exercised.
Changes in `superset/sql/execution/executor.py`:
- `_decode_async_results_payload()` decompresses to bytes and picks JSON or
MessagePack from the payload itself (a JSON object starts with `{`, never the
first byte of a MessagePack map). Results stay readable if the setting differs
between the worker that wrote them and the process reading them.
- `_decode_async_statement_data()` reads Arrow IPC statement data with
`SupersetResultSet.convert_table_to_df()` (the conversion SQL Lab uses),
returns `None` for DML, and keeps the columns of empty results.
- `_records_to_dataframe()` builds DataFrames from JSON records through
Arrow, so an integer column with nulls keeps `None`/`1` like the MessagePack
and sync paths, instead of `NaN`/`1.0`. Mixed or nested values stay Python
objects.
- Cache hits: `_store_in_cache()` stores records plus column names, and
`_get_from_cache()` rebuilds a DataFrame from them. Before this, every sync or
async cache hit returned `StatementResult.data` as a list of dicts, which
breaks the `pd.DataFrame | None` contract, and zero-row results lost their
columns. Older cache entries (DataFrames, or records without columns) still
load.
Stopped, failed and timed-out handles are unchanged. They report the stored
status and error and do not read the results backend; new tests cover this.
Related: this was found during acceptance testing of
https://github.com/apache/superset/pull/45283. It is a separate, pre-existing
issue in code that PR does not touch.
### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
Not applicable (backend only).
Before, from a real Celery worker with Postgres 17 and a Redis results
backend:
```
status=failed stmts=0 err="Error loading results: 'utf-8' codec can't decode
byte 0x85 in position 0: invalid start byte" # MessagePack
status=failed stmts=0 err="Error loading results: a bytes-like object is
required, not 'str'" # JSON
```
After:
```
status=success stmts=1 err=None data=[{'value': 42}]
```
### TESTING INSTRUCTIONS
Unit tests: `pytest tests/unit_tests/sql/execution/test_async_results.py`.
These tests write results with the real Celery task writer
(`_finalize_successful_query` + `_store_results_in_backend`) and read them back
with the real handle reader. Nothing in the compression or serialization path
is mocked. They cover MessagePack/Arrow and JSON, a change of the msgpack
setting between writer and reader, single SELECT, zero-row SELECT, nulls and
mixed types, DML, multi-statement scripts, stopped/failed/timed-out handles,
and cache hits. Against the unchanged executor, 14 of these 17 tests fail. The
3 that pass are the stopped/failed/timed-out guards, which already behaved
correctly.
Manual check with real dependencies:
1. Configure Postgres as the metadata/target database, Redis as
`RESULTS_BACKEND`, and a Celery broker. Start a Celery worker.
2. In an app and request context, as an Admin user and as a Gamma user with
database access, run:
```python
handle = database.execute_async(
"SELECT 42 AS value",
QueryOptions(cache=CacheOptions(force_refresh=True)),
)
# poll handle.get_status() until it is no longer pending/running
handle.get_result()
```
3. Expect `SUCCESS` with one statement whose DataFrame contains `42`. Repeat
with `RESULTS_BACKEND_USE_MSGPACK = False`.
I ran this against Postgres 17, Redis 7 and a real Celery worker, with both
`RESULTS_BACKEND_USE_MSGPACK` values. Before the fix, every successful query
failed in both modes. After the fix, all of these pass in both modes:
- `SELECT 42`, 6 times each as Admin and as Gamma
- zero-row SELECT (keeps columns `a`, `b`)
- nulls (`[{'a': None, 'b': 'x'}, {'a': 1, 'b': 'zé'}]`)
- a failing query (reports `FAILED` with the database error)
- a 4-statement DDL/DML/SELECT script
- an async cache hit after a sync execute (`is_cached=True`, DataFrame data)
### ADDITIONAL INFORMATION
- [ ] Has associated issue:
- [ ] Required feature flags:
- [ ] Changes UI
- [ ] Includes DB Migration (follow approval process in
[SIP-59](https://github.com/apache/superset/issues/13351))
- [ ] Migration is atomic, supports rollback & is backwards-compatible
- [ ] Confirm DB migration upgrade and downgrade tested
- [ ] Runtime estimates and downtime expectations provided
- [ ] Introduces new feature or API
- [ ] Removes existing feature or API
--
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]
