fix(export): apply SQL mutation and user impersonation in streaming CSV export - #42412
Conversation
dosubot
Bot
added
authentication:row-level-security
Code Review Agent Run #4babe3Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
|
cc @rusackas @sadpandajoe @villebro @michael-s-molina — this fixes a security/functional regression in the streaming CSV export path introduced by #35478. The path was skipping SQL mutation (causing Trino crashes) and user impersonation (running as service principal instead of the logged-in user). Would appreciate a review. |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #42412 +/- ##
==========================================
- Coverage 65.25% 64.65% -0.61%
==========================================
Files 2795 2796 +1
Lines 157639 157983 +344
Branches 36052 36103 +51
==========================================
- Hits 102869 102143 -726
- Misses 52793 53844 +1051
- Partials 1977 1996 +19
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
47664bb to
dc66e2c
Compare
Code Review Agent Run #4c3afaActionable Suggestions - 0Additional Suggestions - 1
Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
dc66e2c to
8cb3a21
Compare
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
| with closing( | ||
| merged_database.get_raw_connection(catalog=catalog, schema=schema) | ||
| ) as conn: |
There was a problem hiding this comment.
Suggestion: get_raw_connection() is used with closing(), which only calls close() and does not enter or exit a context manager. Since this API provides the connection through a context manager to apply lifecycle handling such as OAuth2 cleanup, SSH tunnel management, and impersonation context, conn becomes the context-manager object rather than the raw DBAPI connection; conn.cursor() therefore fails or bypasses the required setup. Use the connection as a context manager directly and keep cursor cleanup inside that context. [resource leak]
Severity Level: Critical 🚨
- ❌ Streaming CSV exports fail before query execution.
- ❌ User impersonation and OAuth2 setup are bypassed.
- ❌ SSH tunnel lifecycle handling is not entered or exited.
- ⚠️ Export consumers receive stream errors instead of CSV output.Steps of Reproduction ✅
1. Start a streaming CSV export through the streaming export command; `run()` at
`superset/commands/streaming_export/base.py:295` returns a generator whose consumption
executes `_execute_query_and_stream()`.
2. Consume the returned generator so execution reaches
`superset/commands/streaming_export/base.py:243-245`, where
`merged_database.get_raw_connection(catalog=catalog, schema=schema)` is wrapped with
`contextlib.closing` instead of being entered with a `with` statement.
3. `get_raw_connection()` provides a context-managed connection so its `__enter__` logic
establishes the raw connection and its impersonation, OAuth2, and SSH-tunnel contexts;
`closing` does not invoke that context manager's `__enter__`.
4. The assigned `conn` is therefore the context-manager wrapper rather than the DBAPI
connection, and `conn.cursor()` at `superset/commands/streaming_export/base.py:246` fails
before `cursor.execute(sql)` at line 248. The export generator emits a stream failure
instead of CSV data; cleanup is also not performed through the connection context.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/commands/streaming_export/base.py
**Line:** 243:245
**Comment:**
*Resource Leak: `get_raw_connection()` is used with `closing()`, which only calls `close()` and does not enter or exit a context manager. Since this API provides the connection through a context manager to apply lifecycle handling such as OAuth2 cleanup, SSH tunnel management, and impersonation context, `conn` becomes the context-manager object rather than the raw DBAPI connection; `conn.cursor()` therefore fails or bypasses the required setup. Use the connection as a context manager directly and keep cursor cleanup inside that context.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
|
The flagged issue is correct. Using To resolve this, you should use the connection directly as a context manager. Here is the corrected implementation for # Use get_raw_connection() as a context manager directly
with merged_database.get_raw_connection(catalog=catalog, schema=schema) as conn:
cursor = conn.cursor()
try:
cursor.execute(sql)
columns = (
[desc[0] for desc in cursor.description]
if cursor.description
else []
)This ensures the connection's superset/commands/streaming_export/base.py |
Code Review Agent Run #0a2450Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
|
Thanks for catching this. Routing through One thing to confirm before this lands: the previous path set Separately, CI is currently red ( |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Fixes the streaming CSV export path to match non-streaming exports by applying database SQL mutation hooks and ensuring database connections go through the raw-connection context (impersonation/SSH/OAuth2), with accompanying unit test updates and regressions.
Changes:
- Apply
database.mutate_sql_based_on_config(sql)before executing streaming export SQL. - Switch streaming execution from SQLAlchemy engine/
stream_resultsto raw DBAPI cursor viaget_raw_connection(). - Update/add unit tests to validate SQL mutation and raw-connection usage.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| superset/commands/streaming_export/base.py | Applies SQL mutation and executes via raw DBAPI cursor/connection path intended to support impersonation and other connection wrappers. |
| tests/unit_tests/commands/chart/streaming_export_command_test.py | Refactors mocks to DBAPI cursor semantics and adds regression tests for SQL mutation + raw connection usage. |
| with closing( | ||
| merged_database.get_raw_connection(catalog=catalog, schema=schema) | ||
| ) as conn: | ||
| cursor = conn.cursor() |
| datasource.database.get_raw_connection.return_value.__enter__.return_value = ( | ||
| mock_conn | ||
| ) |
| def test_streaming_execution_options_enabled(mocker: MockerFixture) -> None: | ||
| """Test that streaming execution options are enabled.""" | ||
| """Test that get_raw_connection is used for streaming (not get_sqla_engine).""" |
|
Hi @sha174n — thanks for the detailed review! Great catch on the streaming concern. Let me address both points: 1. Server-side cursor / streamingYou're right that we dropped
The 2. CI redI see the issue — one test was referencing Let me push the update. |
|
Hi @sha174n — thanks for the detailed review! Great catch on the streaming concern. Let me address both points: 1. Server-side cursor / streamingYou are right that we dropped
The 2. CI redI see the issue — one test was referencing Let me push the update. |
…SV export The streaming CSV export path (introduced by apache#35478) bypassed two critical steps that all non-streaming paths perform: 1. SQL mutation — `database.mutate_sql_based_on_config()` was never called, leaving trailing semicolons unstripped for engines like Trino that reject them, causing all streaming CSV exports to crash with __STREAM_ERROR__. 2. User impersonation — the path used `get_sqla_engine()` directly instead of `get_raw_connection()`, which bypasses ENGINE_CONTEXT_MANAGER. On databases with `impersonate_user: true` (Trino, Presto, etc.), all streaming CSV exports ran as the service principal, breaking audit trails and potentially bypassing per-user authorization (Ranger, OPA, RLS views). This fix: - Calls `database.mutate_sql_based_on_config(sql)` before execution - Uses `get_raw_connection()` instead of `get_sqla_engine()` directly, which handles SSH tunnels, OAuth2, and user impersonation - Executes via raw DBAPI cursor (same fetchmany() interface) Fixes apache#40465
8cb3a21 to
ec27fa1
Compare
|
Hi @sha174n — I have pushed the update addressing both points:
Thanks for the thorough review! |
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
| # execution. All non-streaming paths go through this — the streaming | ||
| # path was originally skipping it, which left trailing semicolons | ||
| # unstripped for engines like Trino that reject them. | ||
| sql = database.mutate_sql_based_on_config(sql) |
There was a problem hiding this comment.
Suggestion: The SQL mutator receives the original database object before it is merged into the active session. Streaming generators run after the request session may have ended, so a configured SQL_QUERY_MUTATOR that reads a lazy or session-bound database attribute can raise DetachedInstanceError. Merge the database first and invoke the mutator with the merged instance. [stale reference]
Severity Level: Major ⚠️
- ❌ Configured database-aware SQL mutators can fail during exports.
- ⚠️ Streaming response timing exposes detached ORM state.
- ⚠️ Affected SQL mutation and audit customization paths.Steps of Reproduction ✅
1. Start a chart or SQL Lab streaming export through superset/charts/data/api.py:756-765
or superset/sqllab/api.py:417-435. BaseStreamingCSVExportCommand.run() captures the
database and returns a generator at superset/commands/streaming_export/base.py:129-145, so
query execution occurs during response streaming rather than during command construction.
2. Configure SQL_QUERY_MUTATOR with a function that uses the supplied database keyword
argument, which is explicitly passed by Database.mutate_sql_based_on_config() at
superset/models/core.py:787-805.
3. When the generator reaches superset/commands/streaming_export/base.py:225, it invokes
the mutator on the original database object before entering db.session(future=True) and
before session.merge(database) at base.py:227-230.
4. Because the generator runs after the original request work may have released or
detached its ORM session, a mutator that reads a deferred or session-bound Database
attribute can raise DetachedInstanceError. The merged instance is available at base.py:229
and should be used for mutation before opening the connection.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/commands/streaming_export/base.py
**Line:** 225:225
**Comment:**
*Stale Reference: The SQL mutator receives the original database object before it is merged into the active session. Streaming generators run after the request session may have ended, so a configured `SQL_QUERY_MUTATOR` that reads a lazy or session-bound database attribute can raise `DetachedInstanceError`. Merge the database first and invoke the mutator with the merged instance.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix| # This ensures DBAPI drivers (Trino, PostgreSQL, etc.) fetch | ||
| # rows in manageable chunks instead of buffering the entire | ||
| # result set client-side. | ||
| cursor.arraysize = limit |
There was a problem hiding this comment.
Suggestion: limit is None for chart exports and unlimited SQL Lab exports, but DBAPI cursor.arraysize requires an integer. This assignment can raise a driver-specific type error for every unlimited streaming export, and it is unnecessary because _process_rows() already passes self._chunk_size explicitly to fetchmany(). [type error]
Severity Level: Critical 🚨
- ❌ Unlimited chart CSV exports can fail during cursor setup.
- ❌ Unlimited SQL Lab exports can fail during cursor setup.
- ⚠️ Driver behavior becomes database-dependent.Steps of Reproduction ✅
1. Request a chart streaming CSV export through superset/charts/data/api.py:756-765;
StreamingCSVExportCommand._get_row_limit() explicitly returns None for every chart export
at superset/commands/chart/data/streaming_export_command.py:76-83.
2. After entering the raw connection context, execution reaches
BaseStreamingCSVExportCommand._execute_query_and_stream() at
superset/commands/streaming_export/base.py:194.
3. The code assigns None to cursor.arraysize at
superset/commands/streaming_export/base.py:251. DBAPI cursor arraysize is an integer
fetch-size attribute, and drivers may reject None with a TypeError or equivalent driver
error.
4. The subsequent _process_rows() call at base.py:281-283 already invokes
fetchmany(self._chunk_size), so arraysize is not needed to implement the configured batch
size. The export therefore fails for unlimited chart exports before rows are streamed; SQL
Lab exports with select_sql set also return None from
superset/commands/sql_lab/streaming_export_command.py:119-127.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/commands/streaming_export/base.py
**Line:** 251:251
**Comment:**
*Type Error: `limit` is `None` for chart exports and unlimited SQL Lab exports, but DBAPI `cursor.arraysize` requires an integer. This assignment can raise a driver-specific type error for every unlimited streaming export, and it is unnecessary because `_process_rows()` already passes `self._chunk_size` explicitly to `fetchmany()`.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
Code Review Agent Run #5dae28Actionable Suggestions - 0Filtered by Review RulesBito filtered these suggestions based on rules created automatically for your feedback. Manage rules.
Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
amaannawab923
left a comment
amaannawab923
left a comment
There was a problem hiding this comment.
thanks for digging into this, the missing mutate_sql_based_on_config and the impersonation gap are both real and worth fixing.
one thing on the tests before the rest: they wire get_raw_connection.return_value.__enter__.return_value, which assumes the connection is entered with with ... as conn. with the closing() call in base.py the __enter__ never runs, so the cursor _process_rows reads from isn't the configured mock and the row loop doesn't terminate. so the suite really only passes against the direct context-manager form, might be worth double checking the CI run on this one. left inline notes on the details.
| # 3. OAuth2 — get_raw_connection() wraps execution in | ||
| # check_for_oauth2() context. | ||
| with closing( | ||
| merged_database.get_raw_connection(catalog=catalog, schema=schema) |
There was a problem hiding this comment.
get_raw_connection is a @contextmanager (models/core.py), so it's meant to be entered directly: with database.get_raw_connection(...) as conn. wrapping the call in closing() means __enter__ never runs, so the engine setup inside it (the ENGINE_CONTEXT_MANAGER impersonation this is adding, plus oauth2 and ssh tunnels) gets skipped, and conn ends up being the _GeneratorContextManager instead of the dbapi connection. conn.cursor() then raises AttributeError, which the generator catches and turns into __STREAM_ERROR__, so the export would still fail and the impersonation wouldn't take effect.
every other call site enters it directly (e.g. sql/execution/executor.py, models/core.py). probably just:
with merged_database.get_raw_connection(
catalog=catalog, schema=schema
) as conn:
cursor = conn.cursor()get_raw_connection already closes the connection internally, so the outer closing() isn't needed.
| # This ensures DBAPI drivers (Trino, PostgreSQL, etc.) fetch | ||
| # rows in manageable chunks instead of buffering the entire | ||
| # result set client-side. | ||
| cursor.arraysize = limit |
There was a problem hiding this comment.
two issues here: limit is the export row limit and can be None for unlimited, and cursor.arraysize = None isn't valid for dbapi drivers. also _process_rows calls fetchmany(self._chunk_size) with an explicit size, which overrides arraysize, so this line doesn't actually change the fetch batch size. if the intent is to control the batch, self._chunk_size would be the value, otherwise it can probably be dropped.
| with closing( | ||
| merged_database.get_raw_connection(catalog=catalog, schema=schema) | ||
| ) as conn: | ||
| cursor = conn.cursor() |
There was a problem hiding this comment.
worth confirming the streaming behaviour on this path: the old code used execution_options(stream_results=True), which for psycopg2/postgres sets up a server-side cursor. a raw conn.cursor() on psycopg2 buffers the whole result set client-side, so a large postgres export could pull everything into memory here. for trino the dbapi cursor fetches incrementally so it's fine, but since this is the streaming path it'd be good to keep big postgres exports bounded (server-side cursor where supported, or just call out the tradeoff).
|
Hi @amaannawab923 — thanks for catching the mock wiring issue. I've updated the tests to remove the |
| datasource.database.get_raw_connection.return_value = ( | ||
| mock_conn | ||
| ) |
There was a problem hiding this comment.
Suggestion: The test configures get_raw_connection() to return a raw MagicMock, but the real API returns a context manager that must be entered with with. This causes the tests to exercise an artificial contract and allows the production closing(get_raw_connection(...)) regression to pass unnoticed. Configure the mock as an entered context manager and assert that the context is entered and exited. [api mismatch]
Severity Level: Critical 🚨
- ❌ Unit tests miss production connection-entry failures.
- ⚠️ Streaming export regressions can pass CI undetected.
- ⚠️ Connection cleanup behavior is not verified.Steps of Reproduction ✅
1. Run the chart streaming unit tests, whose common setup patches the database at
`tests/unit_tests/commands/chart/streaming_export_command_test.py:27-48`.
2. Each test assigns a raw `mock_conn` directly as `get_raw_connection.return_value`, for
example at `tests/unit_tests/commands/chart/streaming_export_command_test.py:106-111`,
rather than configuring `return_value.__enter__.return_value`.
3. Production `Database.get_raw_connection()` is a context manager at
`superset/models/core.py:675-692`, and the production code is supposed to enter it with
`with ... as conn`.
4. Consequently, these tests exercise a different API contract: they allow `conn.cursor()`
directly and cannot detect the current `closing(get_raw_connection(...))` failure at
`superset/commands/streaming_export/base.py:243-247`. Configure the entered connection
mock and assert context entry and exit.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** tests/unit_tests/commands/chart/streaming_export_command_test.py
**Line:** 109:111
**Comment:**
*Api Mismatch: The test configures `get_raw_connection()` to return a raw `MagicMock`, but the real API returns a context manager that must be entered with `with`. This causes the tests to exercise an artificial contract and allows the production `closing(get_raw_connection(...))` regression to pass unnoticed. Configure the mock as an entered context manager and assert that the context is entered and exited.
Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
Code Review Agent Run #33fb4fActionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
|
Thanks for picking this up. One blocking issue on the mechanism: Database.get_raw_connection is a @contextmanager, so it has to be entered with |
github-actions
Bot
added
the
requires:rebase
SUMMARY
The streaming CSV export path (introduced by #35478) bypassed two critical
steps that all non-streaming export paths perform:
SQL mutation bypass —
database.mutate_sql_based_on_config()was never called,leaving trailing semicolons unstripped for engines like Trino that reject them,
causing all streaming CSV exports to crash with
__STREAM_ERROR__.User impersonation bypass — the path used
get_sqla_engine()directly insteadof
get_raw_connection(), which bypassesENGINE_CONTEXT_MANAGER. On databaseswith
impersonate_user: true(Trino, Presto, etc.), all streaming CSV exportsran as the service principal, breaking audit trails and potentially bypassing
per-user authorization (Ranger, OPA, RLS views).
This fix:
database.mutate_sql_based_on_config(sql)before executionget_raw_connection()instead ofget_sqla_engine()directly, which handlesSSH tunnels, OAuth2, and user impersonation
fetchmany()interface)BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
Before: Streaming CSV exports crash on Trino with
__STREAM_ERROR__: Export failedand run as the service principal.After: Streaming CSV exports succeed and run as the logged-in user (matching the non-streaming Excel export behavior).
Video:
Screencast from 2026-07-25 11-44-56.webm
TESTING INSTRUCTIONS
Unit tests:
Manual testing (local Docker setup):
docker compose up -dhttp://localhost:8888...menu → Download → Export to CSV__STREAM_ERROR__)SELECT * FROM examples.california_housing_data) and click Download CSVManual testing (Trino/Presto with impersonation):
impersonate_user: true...→ Download → Export to CSV__STREAM_ERROR__)ADDITIONAL INFORMATION