codeant-ai-for-open-source[bot] commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3573311931
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -44,12 +46,61 @@
from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.models.sql_lab import Query
+from superset.superset_typing import (
+ OAuth2ClientConfig,
+ OAuth2State,
+)
from superset.utils import json
from superset.utils.core import get_user_agent, QuerySource
+from superset.utils.oauth2 import encode_oauth2_state, generate_code_challenge
if TYPE_CHECKING:
from superset.models.core import Database
+try:
+ from snowflake.connector.errors import DatabaseError
+except ImportError:
+ # Use a distinct sentinel type when snowflake is not installed to avoid
+ # matching unrelated exception types (using `Exception` would be too
broad).
+ class _SnowflakeDatabaseError(Exception):
+ """Sentinel type to stand in for
snowflake.connector.errors.DatabaseError."""
+
+ pass
+
+ DatabaseError = _SnowflakeDatabaseError
+
+
+class CustomSnowflakeAuthErrorMeta(type):
+ """
+ Metaclass whose ``__instancecheck__`` matches Snowflake's invalid/expired
+ OAuth access-token error, so ``CustomSnowflakeAuthError`` can be used as
the
+ ``oauth2_exception`` that triggers the OAuth2 re-auth dance.
+
+ This is only honored via ``isinstance()`` (the path used by
+ ``BaseEngineSpec.needs_oauth2()``); ``except`` clauses do not call
+ ``__instancecheck__``, so it must not be relied on for exception catching.
+ """
+
+ def __instancecheck__(cls, instance: object) -> bool:
+ """
+ Match Snowflake's invalid/expired OAuth token error, whether it arrives
+ wrapped by SQLAlchemy (e.g. ``Engine``-based execution) or as the raw
+ DBAPI exception — ``BaseEngineSpec.execute()`` runs against a bare
+ cursor and never wraps it, so both shapes must be handled here.
+ """
+ orig: object = instance
+ if isinstance(instance, SqlalchemyDatabaseError):
+ orig = cast(SqlalchemyDatabaseError, instance).orig
+
+ return isinstance(orig, DatabaseError) and "Invalid OAuth access
token" in str(
+ orig
+ )
Review Comment:
**Suggestion:** Storing the OAuth access token in the SQLAlchemy URL makes
the engine cache key token-dependent (the cache key includes
`str(sqlalchemy_url)`), so token refreshes create new cached engines
indefinitely and can cause unbounded engine/pool growth. Keep the URL stable
and pass the token via connection arguments (or otherwise exclude it from
engine cache key material) so refreshed tokens do not create new engine cache
entries. [cache]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ SQL Lab queries create new engines on each token refresh.
- ❌ Dashboard queries against Snowflake grow engine cache unbounded.
- ⚠️ Background jobs using Snowflake OAuth may exhaust process memory.
- ⚠️ Connection pooling becomes ineffective due to fragmented engine cache.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Configure a Snowflake database in Superset using `SnowflakeEngineSpec`
(class at
`superset/db_engine_specs/snowflake.py:135`) with `impersonate_user=True`
and an OAuth2
client so that `database.is_oauth2_enabled()` returns True for this database.
2. From SQL Lab, run a query against this Snowflake database so that
`query.database.get_sqla_engine()` is invoked (call site
`superset/sql_lab.py:23-27`),
which enters `Database.get_sqla_engine()` at `superset/models/core.py:455`.
3. Inside `Database.get_sqla_engine()`, `_get_sqla_engine()` is called
(definition around
`superset/models/core.py:654`), which computes an `access_token` via
`get_oauth2_access_token(...)` and then, when `self.impersonate_user` is
True, calls
`self.db_engine_spec.impersonate_user(...)` (lines
`superset/models/core.py:22-29` in the
second snippet), passing `sqlalchemy_url` and `engine_kwargs`.
4. In `SnowflakeEngineSpec.impersonate_user()` (method at
`superset/db_engine_specs/snowflake.py:58-99`), when `user_token` is present
and OAuth2 is
enabled, the URL is mutated to include the token: `url =
url.update_query_dict({"token":
user_token})` (line 97), and both the modified `url` and `engine_kwargs`
(with
`connect_args["authenticator"] = "oauth"`) are returned.
5. Back in `_get_sqla_engine()`, the per-process engine cache key is
computed as
`(self.id, str(sqlalchemy_url), repr(sorted(engine_kwargs.items())))`
(cache-key
construction around `superset/models/core.py:616-621`), so the stringified
URL—including
its `token` query parameter—is part of the cache key.
6. When the OAuth access token is refreshed (for the same user or for
different users),
`SnowflakeEngineSpec.impersonate_user()` again executes
`url.update_query_dict({"token":
user_token})` with the new token value, changing `str(sqlalchemy_url)` while
`self.id` and
most `engine_kwargs` stay the same.
7. Because `str(sqlalchemy_url)` changes on every token refresh,
`_ENGINE_CACHE` in
`superset/models/core.py` treats each token value as a distinct entry, calls
`create_engine(sqlalchemy_url, **engine_kwargs)` for each new token, and
stores the new
engine without ever evicting the old one, leading to unbounded growth of
cached engines
for a single Snowflake database under normal interactive query activity.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=c1baa1b396cc447cbf9078f8a4f9e1fa&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=c1baa1b396cc447cbf9078f8a4f9e1fa&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/db_engine_specs/snowflake.py
**Line:** 97:97
**Comment:**
*Cache: Storing the OAuth access token in the SQLAlchemy URL makes the
engine cache key token-dependent (the cache key includes
`str(sqlalchemy_url)`), so token refreshes create new cached engines
indefinitely and can cause unbounded engine/pool growth. Keep the URL stable
and pass the token via connection arguments (or otherwise exclude it from
engine cache key material) so refreshed tokens do not create new engine cache
entries.
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%2F36856&comment_hash=a1be5120c4dfc07eda519aed281de566afddcaacda64ee0f56853c2b146c26d9&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=a1be5120c4dfc07eda519aed281de566afddcaacda64ee0f56853c2b146c26d9&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]
