amaannawab923 commented on code in PR #42412:
URL: https://github.com/apache/superset/pull/42412#discussion_r3663942988
##########
superset/commands/streaming_export/base.py:
##########
@@ -219,19 +218,44 @@ 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)
Review 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:
```python
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.
##########
superset/commands/streaming_export/base.py:
##########
@@ -219,19 +218,44 @@ 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:
+ cursor = conn.cursor()
+ # Set cursor.arraysize to control the batch size for
fetchmany().
+ # This ensures DBAPI drivers (Trino, PostgreSQL, etc.) fetch
+ # rows in manageable chunks instead of buffering the entire
+ # result set client-side.
+ cursor.arraysize = limit
Review 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.
##########
superset/commands/streaming_export/base.py:
##########
@@ -219,19 +218,44 @@ 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:
+ cursor = conn.cursor()
Review 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).
--
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]
