rusackas commented on code in PR #42412:
URL: https://github.com/apache/superset/pull/42412#discussion_r4031256333
##########
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:
Good catch, fixed! get_raw_connection() is already a context manager so
wrapping it in closing() skips __enter__, and conn ends up being the
context-manager object instead of the DBAPI connection. Switched to entering it
directly with `with`.
##########
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:
+ cursor = conn.cursor()
Review Comment:
Same issue as Codeant's thread above, fixed in the same push. Entering
get_raw_connection() directly with `with` now instead of wrapping it in
closing().
##########
tests/unit_tests/commands/chart/streaming_export_command_test.py:
##########
@@ -199,12 +185,12 @@ def test_streaming_with_null_values(mocker:
MockerFixture) -> None:
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)."""
Review Comment:
Good catch, renamed to test_streaming_uses_get_raw_connection to match what
it's actually asserting.
##########
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)
Review Comment:
Good catch, fixed. Moved the mutator call to use the merged database instead
of the original, since that's the instance actually guaranteed to still be
attached to the session when this runs.
##########
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:
Good catch, dropped that line. _process_rows already passes chunk_size
explicitly to fetchmany(), so arraysize wasn't doing anything, and arraysize =
None would've broken it for every unlimited export anyway.
##########
tests/unit_tests/commands/chart/streaming_export_command_test.py:
##########
@@ -103,17 +103,11 @@ def test_csv_generation_with_small_dataset(mocker:
MockerFixture) -> None:
[],
]
- mock_connection = mocker.MagicMock()
- mock_connection.execution_options.return_value.execute.return_value = (
- mock_result_proxy
- )
- mock_connection.__enter__.return_value = mock_connection
- mock_connection.__exit__.return_value = None
+ mock_conn = mocker.MagicMock()
+ mock_conn.cursor.return_value = mock_cursor
- mock_engine = mocker.MagicMock()
- mock_engine.connect.return_value = mock_connection
- datasource.database.get_sqla_engine.return_value.__enter__.return_value = (
- mock_engine
+ datasource.database.get_raw_connection.return_value = (
+ mock_conn
)
Review Comment:
Fixed, restored the __enter__-based mock here. Looks like an earlier commit
on this branch had weakened the mock to match the closing() bug instead of
catching it, which is exactly what let this ship.
--
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]
