betodealmeida commented on code in PR #40194:
URL: https://github.com/apache/superset/pull/40194#discussion_r3254706860
##########
superset/commands/streaming_export/base.py:
##########
@@ -160,9 +165,20 @@ def _execute_query_and_stream(
# Merge database to prevent DetachedInstanceError
merged_database = session.merge(database)
- # Execute query with streaming
- with merged_database.get_sqla_engine() as engine:
+ with merged_database.get_sqla_engine(
+ catalog=catalog, schema=schema
+ ) as engine:
with engine.connect() as connection:
+ # Run prequeries (e.g. SET search_path for PostgreSQL)
before
+ # streaming. get_prequeries() is the only mechanism that
sets
+ # search_path for PostgreSQL — adjust_engine_params()
ignores
+ # schema for that engine.
+ for prequery in
merged_database.db_engine_spec.get_prequeries(
+ database=merged_database,
+ catalog=catalog,
+ schema=schema,
+ ):
+ connection.execute(text(prequery))
result_proxy = connection.execution_options(
stream_results=True
Review Comment:
I think we should just use `Database.get_raw_connection()` here, so we don't
have to constantly duplicate logic that is critical. It seems that there are a
few other places where we're calling `engine.connect()` without running the
prequeries.
A better way might be to modify `Database.get_sqla_engine()` to add a hook
to the engine so it automatically runs the prequeries on connect — it might be
the safer approach, since engine creation always uses that method (and will be
even more centralized with https://github.com/apache/superset/pull/34826).
Can you take a stab at moving the prequeries to an event listener? Something
like this:
```python
@event.listens_for(engine, "connect")
def run_prequeries(dbapi_conn, connection_record) -> ...
```
Then we're good as long as we always go through `Database.get_sqla_engine()`.
--
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]
