geido opened a new pull request, #45952: URL: https://github.com/apache/superset/pull/45952
### SUMMARY Streaming CSV exports still fail on databases that use per-user OAuth2, Snowflake in particular. This covers exports with a row limit at or above `CSV_STREAMING_ROW_THRESHOLD`, and SQL Lab streaming downloads. The file the user gets holds only `__STREAM_ERROR__:Export failed. Please try again in some time.` #44425 went after this by resolving `g.user` before the generator runs. On master there are two more problems, and the export only works once both are fixed: 1. **The acting user is detached by the time the body streams.** The WSGI server iterates the response body after the view has returned. By then the default `DBEventLogger` has committed, which expires everything in the session, and the request teardown has removed the metadata session. So the `User` that #44425 hands to the generator is detached and expired, and `Database._get_sqla_engine` raises `DetachedInstanceError` as soon as it reads `g.user.id` to look up the OAuth2 token. 2. **The generator has no request context.** It replayed `flask.g` inside a fresh app context, and that can't bring the request context back. `SnowflakeEngineSpec.is_oauth2_enabled()` and `impersonate_user()` turn OAuth2 off outside a request on purpose (to keep it away from alerts and reports). So even with a usable user, the connection goes out without the token and fails with `251006: Password is empty`. The OAuth2 redirect URI (when it falls back to `url_for`) and the query source given to `DB_CONNECTION_MUTATOR` come from the request too. The change, all in `BaseStreamingCSVExportCommand`: - **Run inside the request again.** `run()` records the request's app and request contexts in the view and pushes them again only while the body is iterated. This is what `flask.stream_with_context` does from Flask 3.1.2. I didn't call `stream_with_context` itself, because earlier releases (3.1.1 and before, which `pyproject.toml` still allows) push the contexts as soon as the body is created. If the response then never reaches the server (an `after_request` hook fails, for example), those contexts stay on the worker thread, and later requests on that thread inherit the exporting request's `g` and logged-in user. This replaces the `flask.g` snapshot and replay (`capture_g_context` / `preserve_g_context`). - **Reload the acting user if teardown detached it.** The generator loads it by primary key and points `g.user` (and Flask-Login's cached user) at the fresh instance. It reloads rather than merges so values read during the request can't be written back by a later commit. If the user no longer exists it fails closed. Users that aren't ORM instances, such as embedded guest users, are left alone. - **Commit right after the engine is built.** Building the engine can refresh the user's OAuth2 token, and that refresh is only staged on the session. Without the commit it was rolled back when the session closed (always, with Redis-based coordination), which spends a single-use refresh token. The commit also means no metadata transaction stays open while rows stream. I checked the two halves of the context fix one at a time against a running app (Flask's test client also iterates the body after the view returns). With only the request context, the export fails with the `DetachedInstanceError`. With only the user reload, it fails with `Password is empty`. With both, the CSV is complete. Notes: - Request teardown handlers run twice for a streamed response: once when the view returns and once when the body has been sent. This matches Flask 3.1's `stream_with_context`. The registered handlers (Flask-SQLAlchemy's session removal, Flask-Limiter's context reset) are idempotent. - Not covered here, same as in #44425: if the user has no usable token, the streaming path still can't send them through the OAuth2 dance, because the response headers are already out when the connection fails. Streaming also doesn't use `_open_with_oauth2_retry`, so a token the database rejects while the store still considers it valid isn't refreshed. I'll pick both up separately. ### BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF Recorded locally against a Snowflake stand-in: fakesnow's implementation of Snowflake's REST protocol, so the real connector talks to it unchanged, wrapped to accept only OAuth logins with a token from its own OAuth2 endpoints (like a Snowflake user with no password). The table has 150,000 rows, and a 100,000 row limit makes the exports stream. Before: authorization works and the 100k-row query runs with the user's token, but the SQL Lab download and the chart's "Export to .CSV" both fail. The body streamed by `/api/v1/chart/data` is only the error marker. https://github.com/user-attachments/assets/c2f7c97b-0a4b-4c99-b20d-3624460042a2 After: both exports finish, and each downloaded file has the header plus 100,000 rows. https://github.com/user-attachments/assets/2959f8de-80c2-4b7b-8d16-32236797a661 ### TESTING INSTRUCTIONS ```bash pytest tests/unit_tests/commands/chart/streaming_export_command_test.py \ tests/unit_tests/commands/sql_lab/streaming_export_command_test.py ``` The new tests cover: - the export running in its request context after the view has returned; - a detached acting user being reloaded (Flask-Login's `current_user` included), and failing closed when the user is gone; - guest, attached and missing users being left alone; - the token commit happening before rows stream; - no request context being left on the thread when the body is never iterated. The first two fail on master. Each request in these tests pushes its own app context, as `Flask.wsgi_app` does, so they can't pass for the wrong reason. Manually: 1. Set up a Snowflake database with OAuth2 (`oauth2_client_info` on the connection, or `DATABASE_OAUTH2_CLIENTS`) and "Impersonate logged in user" turned on. 2. Authorize as a regular user. 3. Export a table chart with a row limit of 100,000 to CSV, and download the results of a SQL Lab query run with a limit of 100,000. Before: "Export failed. Please try again in some time." After: the full CSV. ### ADDITIONAL INFORMATION - [x] Has associated issue: follow-up to #44425 - [ ] 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 🤖 Generated with [Claude Code](https://claude.com/claude-code) -- 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]
