codeant-ai-for-open-source[bot] commented on code in PR #40194:
URL: https://github.com/apache/superset/pull/40194#discussion_r3255645435
##########
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:
**🟠Architect Review — HIGH**
get_sqla_engine() now registers a SQLAlchemy "connect" event that runs
db_engine_spec.get_prequeries(), but get_raw_connection() still calls
get_prequeries() and executes the returned statements on the same raw
connection, so prequeries run twice for all get_raw_connection() callers. For
engines with non-idempotent prequeries (e.g.
StarRocksEngineSpec.get_prequeries() issuing `EXECUTE AS ... WITH NO REVERT;`),
this can break queries or impersonation by re-running the same pre-session
command on every new raw connection.
**Suggestion:** Make prequery execution single-sourced: either rely solely
on the engine-level connect listener (and remove the manual get_prequeries()
loop from get_raw_connection()), or add a guard so that get_raw_connection()
does not re-run prequeries when using get_sqla_engine(). Add a regression test
that asserts prequeries are executed exactly once per connection for a
get_raw_connection() path, including a case using StarRocks impersonation.
[Fix in
Cursor](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt=This%20is%20an%20%2A%2AArchitect%20%2F%20Logical%20Review%2A%2A%20comment%20left%20during%20a%20code%20review.%20These%20reviews%20are%20first-class%2C%20important%20findings%20%E2%80%94%20not%20optional%20suggestions.%20Do%20NOT%20dismiss%20this%20as%20a%20%27big%20architectural%20change%27%20just%20because%20the%20title%20says%20architect%20review%3B%20most%20of%20these%20can%20be%20resolved%20with%20a%20small%2C%20localized%20fix%20once%20the%20intent%20is%20understood.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmodels%2Fcore.py%0A%2A%2ALine%3A%2A%2A%20478%3A495%0A%2A%2AComment%3A%2A%2A%0A%09%2AHIGH%3A%20get_sqla_engine%28%29%20now%20registers%20a%20SQLAlchemy%20%22connect%22%20event%20that%20runs%20db_engine_spec.get_prequeries%28%29%2C%20but%20get_raw_connection%28%29%20still%20calls%20get_prequeries%28%29%20and%20executes%20the%20returned%20statements%20on%20the%20same%20raw%20connection%2C%20so%20prequeries%20r
un%20twice%20for%20all%20get_raw_connection%28%29%20callers.%20For%20engines%20with%20non-idempotent%20prequeries%20%28e.g.%20StarRocksEngineSpec.get_prequeries%28%29%20issuing%20%60EXECUTE%20AS%20...%20WITH%20NO%20REVERT%3B%60%29%2C%20this%20can%20break%20queries%20or%20impersonation%20by%20re-running%20the%20same%20pre-session%20command%20on%20every%20new%20raw%20connection.%0A%0AValidate%20the%20correctness%20of%20the%20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AIf%20a%20suggested%20approach%20is%20provided%20above%2C%20use%20it%20as%20the%20authoritative%20instruction.%20If%20no%20explicit%20code%20suggestion%20is%20given%2C%20you%20MUST%20still%20draft%20and%20apply%20your%20own%20minimal%2C%20localized%20fix%20%E2%80%94%20do%20not%20punt%20back%20with%20%27no%20suggestion%20provided%2C%20review%20manually%27.%20Keep%20the%20change%20as%20small%20as%20possible%
3A%20add%20a%20guard%20clause%2C%20gate%20on%20a%20loading%20state%2C%20reorder%20an%20await%2C%20wrap%20in%20a%20conditional%2C%20etc.%20Do%20not%20refactor%20surrounding%20code%20or%20expand%20scope%20beyond%20the%20finding.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
| [Fix in VSCode
Claude](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt=This%20is%20an%20%2A%2AArchitect%20%2F%20Logical%20Review%2A%2A%20comment%20left%20during%20a%20code%20review.%20These%20reviews%20are%20first-class%2C%20important%20findings%20%E2%80%94%20not%20optional%20suggestions.%20Do%20NOT%20dismiss%20this%20as%20a%20%27big%20architectural%20change%27%20just%20because%20the%20title%20says%20architect%20review%3B
%20most%20of%20these%20can%20be%20resolved%20with%20a%20small%2C%20localized%20fix%20once%20the%20intent%20is%20understood.%0A%0A%2A%2APath%3A%2A%2A%20superset%2Fmodels%2Fcore.py%0A%2A%2ALine%3A%2A%2A%20478%3A495%0A%2A%2AComment%3A%2A%2A%0A%09%2AHIGH%3A%20get_sqla_engine%28%29%20now%20registers%20a%20SQLAlchemy%20%22connect%22%20event%20that%20runs%20db_engine_spec.get_prequeries%28%29%2C%20but%20get_raw_connection%28%29%20still%20calls%20get_prequeries%28%29%20and%20executes%20the%20returned%20statements%20on%20the%20same%20raw%20connection%2C%20so%20prequeries%20run%20twice%20for%20all%20get_raw_connection%28%29%20callers.%20For%20engines%20with%20non-idempotent%20prequeries%20%28e.g.%20StarRocksEngineSpec.get_prequeries%28%29%20issuing%20%60EXECUTE%20AS%20...%20WITH%20NO%20REVERT%3B%60%29%2C%20this%20can%20break%20queries%20or%20impersonation%20by%20re-running%20the%20same%20pre-session%20command%20on%20every%20new%20raw%20connection.%0A%0AValidate%20the%20correctness%20of%20the%
20flagged%20issue.%20If%20correct%2C%20How%20can%20I%20resolve%20this%3F%20If%20you%20propose%20a%20fix%2C%20implement%20it%20and%20please%20make%20it%20concise.%0AIf%20a%20suggested%20approach%20is%20provided%20above%2C%20use%20it%20as%20the%20authoritative%20instruction.%20If%20no%20explicit%20code%20suggestion%20is%20given%2C%20you%20MUST%20still%20draft%20and%20apply%20your%20own%20minimal%2C%20localized%20fix%20%E2%80%94%20do%20not%20punt%20back%20with%20%27no%20suggestion%20provided%2C%20review%20manually%27.%20Keep%20the%20change%20as%20small%20as%20possible%3A%20add%20a%20guard%20clause%2C%20gate%20on%20a%20loading%20state%2C%20reorder%20an%20await%2C%20wrap%20in%20a%20conditional%2C%20etc.%20Do%20not%20refactor%20surrounding%20code%20or%20expand%20scope%20beyond%20the%20finding.%0AOnce%20fix%20is%20implemented%2C%20also%20check%20other%20comments%20on%20the%20same%20PR%2C%20and%20ask%20user%20if%20the%20user%20wants%20to%20fix%20the%20rest%20of%20the%20comments%20as%20well.
%20if%20said%20yes%2C%20then%20fetch%20all%20the%20comments%20validate%20the%20correctness%20and%20implement%20a%20minimal%20fix%0A)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is an **Architect / Logical Review** comment left during a code review.
These reviews are first-class, important findings — not optional suggestions.
Do NOT dismiss this as a 'big architectural change' just because the title says
architect review; most of these can be resolved with a small, localized fix
once the intent is understood.
**Path:** superset/models/core.py
**Line:** 478:495
**Comment:**
*HIGH: get_sqla_engine() now registers a SQLAlchemy "connect" event
that runs db_engine_spec.get_prequeries(), but get_raw_connection() still calls
get_prequeries() and executes the returned statements on the same raw
connection, so prequeries run twice for all get_raw_connection() callers. For
engines with non-idempotent prequeries (e.g.
StarRocksEngineSpec.get_prequeries() issuing `EXECUTE AS ... WITH NO REVERT;`),
this can break queries or impersonation by re-running the same pre-session
command on every new raw connection.
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.
If a suggested approach is provided above, use it as the authoritative
instruction. If no explicit code suggestion is given, you MUST still draft and
apply your own minimal, localized fix — do not punt back with 'no suggestion
provided, review manually'. Keep the change as small as possible: add a guard
clause, gate on a loading state, reorder an await, wrap in a conditional, etc.
Do not refactor surrounding code or expand scope beyond the finding.
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
```
</details>
--
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]
