fitzee commented on code in PR #40194:
URL: https://github.com/apache/superset/pull/40194#discussion_r3255723568
##########
superset/models/core.py:
##########
@@ -468,13 +468,31 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
engine_context_manager = app.config["ENGINE_CONTEXT_MANAGER"]
with engine_context_manager(self, catalog, schema):
with check_for_oauth2(self):
- yield self._get_sqla_engine(
+ engine = self._get_sqla_engine(
catalog=catalog,
schema=schema,
nullpool=nullpool,
source=source,
sqlalchemy_uri=sqlalchemy_uri,
)
+ prequeries = self.db_engine_spec.get_prequeries(
+ database=self,
+ catalog=catalog,
+ schema=schema,
+ )
+ if prequeries:
+
+ def run_prequeries(
+ dbapi_connection: Any,
+ connection_record: Any, # pylint:
disable=unused-argument
+ ) -> None:
+ cursor = dbapi_connection.cursor()
+ for prequery in prequeries:
+ cursor.execute(prequery)
+ cursor.close()
+
+ sqla.event.listen(engine, "connect", run_prequeries)
+ yield engine
Review Comment:
Fixed in commit
[0ceaf8ff11](https://github.com/fitzee/superset/commit/0ceaf8ff11390001db6f0b49a8fd13b331280857):
the manual `get_prequeries()` loop was removed from `get_raw_connection()`
entirely. The connect event listener registered by `get_sqla_engine()` is now
the sole execution point, so prequeries fire exactly once regardless of which
call path is used.
Regression test added in commit
[97fd0e4882](https://github.com/fitzee/superset/commit/97fd0e4882e9f35a5b9afb9ae2ceade4bb2e07c8):
`test_get_raw_connection_executes_prequeries_exactly_once` intercepts
`sqla.event.listen`, simulates SQLAlchemy firing the connect event when
`raw_connection()` is called, and asserts `cursor.execute` is called exactly
once — not twice, not zero.
##########
superset/models/core.py:
##########
@@ -468,13 +468,31 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
engine_context_manager = app.config["ENGINE_CONTEXT_MANAGER"]
with engine_context_manager(self, catalog, schema):
with check_for_oauth2(self):
- yield self._get_sqla_engine(
+ engine = self._get_sqla_engine(
catalog=catalog,
schema=schema,
nullpool=nullpool,
source=source,
sqlalchemy_uri=sqlalchemy_uri,
)
+ prequeries = self.db_engine_spec.get_prequeries(
+ database=self,
+ catalog=catalog,
+ schema=schema,
+ )
+ if prequeries:
+
+ def run_prequeries(
+ dbapi_connection: Any,
+ connection_record: Any, # pylint:
disable=unused-argument
+ ) -> None:
+ cursor = dbapi_connection.cursor()
+ for prequery in prequeries:
+ cursor.execute(prequery)
+ cursor.close()
Review Comment:
Fixed in commit
[c9b1d64dee](https://github.com/fitzee/superset/commit/c9b1d64dee9b1f0fb7d90543f9d3bcd17b9d0447):
`run_prequeries` now wraps cursor usage in `try/finally`, so `cursor.close()`
is called unconditionally even when a prequery raises.
Test added in the same commit:
`test_get_sqla_engine_prequery_cursor_closed_on_exception` injects an exception
from `cursor.execute` and asserts `cursor.close` is still called exactly once.
--
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]
