codeant-ai-for-open-source[bot] commented on code in PR #42412:
URL: https://github.com/apache/superset/pull/42412#discussion_r3654134542
##########
superset/commands/streaming_export/base.py:
##########
@@ -219,19 +218,39 @@ def _execute_query_and_stream(
delimiter = csv_export_config.get("sep", ",")
decimal_separator = csv_export_config.get("decimal", ".")
+ # Apply SQL mutations (e.g. SQL_QUERY_MUTATOR config hook) before
+ # 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)
+
with db.session(future=True) as session:
# Merge database to prevent DetachedInstanceError
merged_database = session.merge(database)
- with merged_database.get_sqla_engine(
- catalog=catalog, schema=schema
- ) as engine:
- with engine.connect() as connection:
- result_proxy = connection.execution_options(
- stream_results=True
- ).execute(text(sql))
-
- columns = list(result_proxy.keys())
+ # Use get_raw_connection() instead of get_sqla_engine() directly.
+ # This is critical for:
+ # 1. User impersonation — get_raw_connection() goes through the
+ # ENGINE_CONTEXT_MANAGER which applies impersonate_user settings
+ # (e.g. X-Trino-User header). Without this, all streaming CSV
+ # exports run as the service principal, breaking audit trails
+ # and potentially bypassing per-user authorization (Ranger, OPA,
+ # RLS views).
+ # 2. SSH tunnels — get_raw_connection() sets up SSH tunnels if
+ # configured on the database.
+ # 3. OAuth2 — get_raw_connection() wraps execution in
+ # check_for_oauth2() context.
+ with closing(
+ merged_database.get_raw_connection(catalog=catalog,
schema=schema)
+ ) as conn:
Review 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]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ 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.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
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.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=4dd8f94a555b43559cfc8a7982473a91&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=4dd8f94a555b43559cfc8a7982473a91&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
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
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42412&comment_hash=3ce53f2385557ea58253fa631c35e548a0076122d808427138733aa47aea0324&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42412&comment_hash=3ce53f2385557ea58253fa631c35e548a0076122d808427138733aa47aea0324&reaction=dislike'>👎</a>
--
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]
