fix(database): move prequeries from event listeners to direct execution#42079
fix(database): move prequeries from event listeners to direct execution#42079alseerx wants to merge 5 commits into
Conversation
Removed prequery handling from engine context management and added it to the raw connection context.
Refactor tests to check that get_sqla_engine does not register event listeners for prequeries and that get_raw_connection executes prequeries directly on the connection.
dosubot
Bot
added
change:backend
![codeant-ai-for-open-source[bot]](https://avatars.githubusercontent.com/u/151810221?s=60&v=4){kind=small}
| prequeries = self.db_engine_spec.get_prequeries( | ||
| database=self, | ||
| catalog=catalog, | ||
| schema=schema, |
There was a problem hiding this comment.
Suggestion: Add an explicit type annotation for this newly introduced local variable to comply with the project’s type-hinting requirement. [custom_rule]
Severity Level: Minor 🧹
Why it matters? ⭐
This is newly added Python code that introduces a local variable without a type annotation. Since the project rule requires type hints on relevant annotatable variables, this is a real violation.
Rule source 📖
.cursor/rules/dev-standard.mdc (line 28)
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 648:651
**Comment:**
*Custom Rule: Add an explicit type annotation for this newly introduced local variable to comply with the project’s type-hinting requirement.
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| schema=schema, | ||
| ) | ||
| if prequeries: | ||
| cursor = conn.cursor() |
There was a problem hiding this comment.
Suggestion: Add a type annotation for this cursor variable so the new database execution block remains fully type-hinted. [custom_rule]
Severity Level: Minor 🧹
Why it matters? ⭐
This is a newly added local variable in Python code and it is not type-annotated. That matches the type-hinting rule for new or modified code, so the suggestion is valid.
Rule source 📖
.cursor/rules/dev-standard.mdc (line 28)
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 654:654
**Comment:**
*Custom Rule: Add a type annotation for this cursor variable so the new database execution block remains fully type-hinted.
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| database = Database(database_name="my_db", sqlalchemy_uri="postgresql://") | ||
| with database.get_sqla_engine(catalog=None, schema="bad_schema"): | ||
| pass | ||
| prequery = 'SET search_path = "my_schema"' |
There was a problem hiding this comment.
Suggestion: Add an explicit type annotation for this newly introduced local variable so the new code fully complies with the type-hinting requirement. [custom_rule]
Severity Level: Minor 🧹
Why it matters? ⭐
The new local variable is introduced without a type annotation in Python code, and it can clearly be annotated as a string. This matches the rule requiring type hints on relevant variables that can be annotated.
Rule source 📖
.cursor/rules/dev-standard.mdc (line 28)
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** tests/unit_tests/models/core_test.py
**Line:** 753:753
**Comment:**
*Custom Rule: Add an explicit type annotation for this newly introduced local variable so the new code fully complies with the type-hinting requirement.
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![bito-code-review[bot]](https://avatars.githubusercontent.com/in/1061978?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review Agent Run #b1c9c6
Actionable Suggestions - 1
-
superset/models/core.py - 1
- Runtime NameError for undefined variable · Line 506-506
Review Details
-
Files reviewed - 2 · Commit Range:
95b93ce..050c997- superset/models/core.py
- tests/unit_tests/models/core_test.py
-
Files skipped - 0
-
Tools
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
Bito Usage Guide
Commands
Type the following command in the pull request comment and save the comment.
-
/review- Manually triggers a full AI review. -
/pause- Pauses automatic reviews on this pull request. -
/resume- Resumes automatic reviews. -
/resolve- Marks all Bito-posted review comments as resolved. -
/abort- Cancels all in-progress reviews.
Refer to the documentation for additional commands.
Configuration
This repository uses Superset You can customize the agent settings here or contact your Bito workspace admin at evan@preset.io.
Documentation & Help
AI Code Review powered by
| @@ -518,35 +505,7 @@ def get_sqla_engine( # pylint: disable=too-many-arguments | |||
| sqlalchemy_uri=sqlalchemy_uri, | |||
| cacheable=not prequeries, | |||
There was a problem hiding this comment.
Line 506 references undefined variable prequeries which will raise NameError at runtime. The prequery event-listener logic was removed from _get_sqla_engine() and moved to get_raw_connection() (lines 648-659), but the reference at line 506 was left behind. Remove the cacheable=not prequeries argument.
Code Review Run #b1c9c6
Should Bito avoid suggestions like this for future reviews? (Manage Rules)
- Yes, avoid them
| prequeries = self.db_engine_spec.get_prequeries( | ||
| database=self, | ||
| catalog=catalog, | ||
| schema=schema, | ||
| ) | ||
| if prequeries: | ||
| cursor = conn.cursor() | ||
| try: | ||
| for prequery in prequeries: | ||
| cursor.execute(prequery) | ||
| finally: | ||
| cursor.close() | ||
| yield conn |
There was a problem hiding this comment.
Suggestion: This change executes prequeries only inside get_raw_connection(), which means code paths that use get_sqla_engine() + engine.connect() no longer receive required pre-session statements. That breaks schema/catalog/session setup for engines that rely on get_prequeries (for example Postgres set search_path) and can run queries against the wrong schema. Apply the same prequery behavior to SQLAlchemy engine.connect() call paths (or centralize it so both connection acquisition paths are covered). [incomplete implementation]
Severity Level: Major ⚠️
⚠️ SQL Lab cancellation connections omit dynamic-schema prequeries.
⚠️ Engines using schema-dependent prequeries may use default schema.
⚠️ User impersonation via get_prequeries may not apply everywhere.Steps of Reproduction ✅
1. Note that db_engine_specs such as Postgres, DB2, StarRocks implement get_prequeries()
(e.g. superset/db_engine_specs/postgres.py:717–23) to return pre-session statements like
`SET search_path = "my_schema"` for dynamic schema and permission enforcement (documented
in superset/db_engine_specs/base.py:1690–39).
2. In the refactored Database.get_raw_connection() in superset/models/core.py:633–211,
prequeries are fetched and executed only inside the raw DBAPI connection path: `prequeries
= self.db_engine_spec.get_prequeries(...); for prequery in prequeries:
cursor.execute(prequery)` at lines 199–208, then `yield conn` at line 211.
3. Several real call sites still obtain connections via get_sqla_engine rather than
get_raw_connection, for example SQL Lab cancellation in superset/sql_lab.py:33–40 which
uses `with query.database.get_sqla_engine(catalog=query.catalog, schema=query.schema) as
engine:` and then `engine.raw_connection()` directly, bypassing get_raw_connection’s
prequery execution.
4. Because prequeries are not executed anywhere in get_sqla_engine() and
engine.connect()/engine.raw_connection() paths that do not go through get_raw_connection,
connections used there will miss required pre-session statements (schema search_path,
impersonation EXECUTE AS, catalog/schema selection), so queries or operations executed via
those connections can run under the wrong schema or user context, diverging from the
dynamic-schema semantics described in db_engine_specs.base.get_prequeries.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 648:660
**Comment:**
*Incomplete Implementation: This change executes prequeries only inside `get_raw_connection()`, which means code paths that use `get_sqla_engine()` + `engine.connect()` no longer receive required pre-session statements. That breaks schema/catalog/session setup for engines that rely on `get_prequeries` (for example Postgres `set search_path`) and can run queries against the wrong schema. Apply the same prequery behavior to SQLAlchemy `engine.connect()` call paths (or centralize it so both connection acquisition paths are covered).
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 fixFix indentation and formatting issues in core.py.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #42079 +/- ##
==========================================
- Coverage 65.08% 65.02% -0.07%
==========================================
Files 2747 2751 +4
Lines 153905 154376 +471
Branches 35293 35373 +80
==========================================
+ Hits 100164 100377 +213
- Misses 51830 52087 +257
- Partials 1911 1912 +1
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Code Review Agent Run #c8e218Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
✅ Deploy Preview for superset-docs-preview ready!
To edit notification comments on pull requests, go to your Netlify project configuration. |
| prequeries = self.db_engine_spec.get_prequeries( | ||
| database=self, | ||
| catalog=catalog, | ||
| schema=schema, |
There was a problem hiding this comment.
Suggestion: Add an explicit type annotation for the prequeries collection to satisfy the type-hint requirement for newly introduced local variables. [custom_rule]
Severity Level: Minor 🧹
Why it matters? ⭐
The new local variable prequeries is introduced without an explicit type annotation, and it is subsequently used in a truthiness check and loop. Under the type-hint rule, this is a newly added, annotatable variable that should be typed.
Rule source 📖
.cursor/rules/dev-standard.mdc (line 28)
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 647:650
**Comment:**
*Custom Rule: Add an explicit type annotation for the prequeries collection to satisfy the type-hint requirement for newly introduced local variables.
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
| prequeries = self.db_engine_spec.get_prequeries( | ||
| database=self, | ||
| catalog=catalog, | ||
| schema=schema, | ||
| ) | ||
| if prequeries: | ||
| cursor = conn.cursor() | ||
| try: | ||
| for prequery in prequeries: | ||
| cursor.execute(prequery) | ||
| finally: | ||
| cursor.close() |
There was a problem hiding this comment.
Suggestion: This change limits get_prequeries execution to get_raw_connection, so any code path that uses get_sqla_engine directly (for example inspector-based metadata calls) no longer runs required pre-session statements. That breaks engine specs that implement impersonation or session setup via prequeries (e.g. StarRocks EXECUTE AS), and can expose metadata using the service account instead of the effective user. Reintroduce prequery application for non-raw-connection engine usage (without dynamic listener mutation), such as applying prequeries on explicit connections used by inspector/metadata paths. [security]
Severity Level: Critical 🚨
- ❌ StarRocks impersonation not applied for metadata inspector connections.
- ❌ Database tables API uses service account for listing tables.
- ⚠️ Streaming export queries ignore prequeries, bypassing user impersonation.
- ⚠️ Schema default prequeries skipped for engine-only metadata operations.Steps of Reproduction ✅
1. Configure a Database row using the StarRocks engine spec with user impersonation
enabled, so that `Database.db_engine_spec` is `StarRocksEngineSpec` and
`database.impersonate_user` is True; see `superset/db_engine_specs/starrocks.py` lines
380–439 where `get_prequeries()` returns an `EXECUTE AS` statement and the comment
explicitly states “User impersonation is actually achieved via `get_prequeries`”.
2. Observe that `Database.get_raw_connection()` in `superset/models/core.py` lines 632–659
calls `self.db_engine_spec.get_prequeries(database=self, catalog=catalog, schema=schema)`
(lines 647–651), then executes each returned prequery on the DBAPI connection cursor
before yielding `conn`, so any SQL execution paths that use `get_raw_connection()` (e.g.
SQL Lab execution paths in `superset/sql/execution/executor.py:538` and
`superset/sql_lab.py:499`) still run the StarRocks `EXECUTE AS` impersonation prequeries.
3. Follow an inspector-based metadata path, for example
`Database.get_all_table_names_in_schema()` in `superset/models/core.py` lines 979–1005,
which calls `self.get_inspector(catalog=catalog, schema=schema)` (line 997);
`get_inspector()` (lines 1085–1091) uses `with self.get_sqla_engine(catalog=catalog,
schema=schema) as engine:` and then `sqla.inspect(engine)`, meaning the Inspector obtains
connections directly from the engine without going through `get_raw_connection()`.
4. Inspect `Database.get_sqla_engine()` in `superset/models/core.py` lines 455–507 and
note that it now only calls `_get_sqla_engine()` and yields the shared engine (line 507)
without invoking `get_prequeries()` or registering any `connect` event listener; a Grep
for `get_prequeries` in the repo shows the only runtime callsite is inside
`get_raw_connection()` (lines 647–651). Therefore, when the database tables API runs
`TableListCommand.run()` in `superset/commands/database/tables.py` lines 60–70, it calls
`self._model.get_all_table_names_in_schema(...)`, which goes through `get_inspector()` and
`get_sqla_engine()` and opens connections where StarRocks `EXECUTE AS` prequeries are
never executed, so metadata inspection and similar engine-only operations run under the
service account instead of the impersonated user.
(Use Cmd/Ctrl + Click for best experience)
Prompt for AI Agent 🤖
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 647:658
**Comment:**
*Security: This change limits `get_prequeries` execution to `get_raw_connection`, so any code path that uses `get_sqla_engine` directly (for example inspector-based metadata calls) no longer runs required pre-session statements. That breaks engine specs that implement impersonation or session setup via prequeries (e.g. StarRocks `EXECUTE AS`), and can expose metadata using the service account instead of the effective user. Reintroduce prequery application for non-raw-connection engine usage (without dynamic listener mutation), such as applying prequeries on explicit connections used by inspector/metadata paths.
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
Code Review Agent Run #30f351Actionable Suggestions - 0Review Details
Bito Usage GuideCommands Type the following command in the pull request comment and save the comment.
Refer to the documentation for additional commands. Configuration This repository uses Documentation & Help AI Code Review powered by |
1 participant
SUMMARY
Fixes the "deque mutated during iteration" error (#40903) caused by dynamically adding/removing SQLAlchemy
connectevent listeners on shared cached engines.Root cause:
get_sqla_engine()registered a temporaryconnectevent listener for prequeries (e.g.SET search_path) on the cached engine, then removed it in afinallyblock. Under concurrent access, multiple threads mutated the engine's internal listener deque simultaneously, triggeringRuntimeError: deque mutated during iteration.Fix: Move prequery execution from event listeners to direct cursor execution in
get_raw_connection(), which runs after obtaining each connection. This eliminates all dynamic listener mutation on shared engines while preserving identical prequery behavior.BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF
N/A — backend-only change, no UI impact.
TESTING INSTRUCTIONS
get_sqla_engine()get_raw_connection()ADDITIONAL INFORMATION
Link to Devin session: https://app.devin.ai/sessions/90ed846778f942dc89855b900484e7f6
Requested by: @alseerx
Credits to original PR: OmarDadabhoy#23
Since @OmarDadabhoy is not responding to requests in his fork to merge the changes to the main repository. I am "adopting" these changes to create the PR to the main repository.