codeant-ai-for-open-source[bot] commented on code in PR #42079:
URL: https://github.com/apache/superset/pull/42079#discussion_r3586106845
##########
superset/models/core.py:
##########
@@ -686,6 +645,18 @@ def get_raw_connection(
) as engine:
with check_for_oauth2(self):
with closing(engine.raw_connection()) as conn:
+ 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
Review 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]
<details>
<summary><b>Severity Level:</b> Major ⚠️</summary>
```mdx
⚠️ 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.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
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.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=4b9212d97dbc4e93b4d296333ddc0415&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=4b9212d97dbc4e93b4d296333ddc0415&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
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 fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42079&comment_hash=275bc975013e9b606da79d62426f5f2e767d46bbed6f7164a7ba72916b1c302f&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42079&comment_hash=275bc975013e9b606da79d62426f5f2e767d46bbed6f7164a7ba72916b1c302f&reaction=dislike'>👎</a>
##########
superset/models/core.py:
##########
@@ -518,35 +505,7 @@ def get_sqla_engine( # pylint: disable=too-many-arguments
sqlalchemy_uri=sqlalchemy_uri,
cacheable=not prequeries,
)
- if prequeries:
- # SQLAlchemy connect event: runs prequeries on every
new
- # DBAPI connection (e.g. SET search_path for
PostgreSQL).
- def run_prequeries(
- dbapi_connection: Any,
- connection_record: Any, # pylint:
disable=unused-argument
- ) -> None:
- cursor = dbapi_connection.cursor()
- try:
- for prequery in prequeries:
- cursor.execute(prequery)
- finally:
- cursor.close()
-
- sqla.event.listen(engine, "connect", run_prequeries)
- try:
- yield engine
- finally:
- sqla.event.remove(engine, "connect",
run_prequeries)
- # The engine is private (cacheable=False above), so
- # nothing else can hold a reference: dispose it to
- # release its pool immediately. With the default
- # nullpool=True this is a no-op safety net; it
- # matters if a caller ever passes nullpool=False,
- # where each private engine would otherwise keep a
- # short-lived QueuePool alive until GC.
- engine.dispose()
- else:
- yield engine
+ yield engine
Review Comment:
**Suggestion:** `get_sqla_engine()` still references `prequeries` when
calling `_get_sqla_engine`, but this variable is no longer defined in the
method after the listener-removal refactor. At runtime this raises `NameError`
before any engine is yielded. Recompute `prequeries` in this scope or remove
this flag entirely now that listener registration was removed. [incorrect
variable usage]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
❌ All get_sqla_engine callers crash with NameError.
❌ Example dataset loaders fail to initialize database engine.
⚠️ SQL Lab cancellation using get_sqla_engine cannot acquire engine.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Trigger any code path that uses the Database.get_sqla_engine() context
manager, for
example load_world_bank_health() in superset/examples/world_bank.py which
calls `with
database.get_sqla_engine() as engine:` at line 54.
2. The call enters get_sqla_engine() in superset/models/core.py
(contextmanager starting
around line 6 of the hunk), eventually executing the assignment `engine =
self._get_sqla_engine(..., cacheable=not prequeries)` at diff lines 500–507.
3. In the current implementation of superset/models/core.py, `prequeries` is
not defined
anywhere in get_sqla_engine() or at module scope (verified via grep; only
usages are this
line and the local variable in get_raw_connection at lines 199–210).
4. When Python evaluates `cacheable=not prequeries`, name resolution fails
and a NameError
is raised before the engine is created, so the context manager never reaches
`yield
engine` (line 508) and the caller’s with-block crashes instead of receiving
a usable
SQLAlchemy Engine.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=3b0365da691e4d5c8a03baea7481fbd5&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=3b0365da691e4d5c8a03baea7481fbd5&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/models/core.py
**Line:** 500:508
**Comment:**
*Incorrect Variable Usage: `get_sqla_engine()` still references
`prequeries` when calling `_get_sqla_engine`, but this variable is no longer
defined in the method after the listener-removal refactor. At runtime this
raises `NameError` before any engine is yielded. Recompute `prequeries` in this
scope or remove this flag entirely now that listener registration was removed.
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
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42079&comment_hash=8726725015b54e3429d64711e16cd709649335fef64ab4cbded48f06a6bcdb75&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42079&comment_hash=8726725015b54e3429d64711e16cd709649335fef64ab4cbded48f06a6bcdb75&reaction=dislike'>👎</a>
--
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]
