codeant-ai-for-open-source[bot] commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3564726709
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -44,12 +46,53 @@
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:
+ return (
+ isinstance(instance, SqlalchemyDatabaseError)
+ and isinstance(cast(SqlalchemyDatabaseError, instance).orig,
DatabaseError)
+ and "Invalid OAuth access token" in str(instance)
+ )
Review Comment:
**Suggestion:** The OAuth error matcher only recognizes SQLAlchemy-wrapped
`DatabaseError` instances with `.orig`, but query execution paths use raw DBAPI
cursor exceptions as well. When token failures surface as direct Snowflake
connector errors, this check returns false and the OAuth re-auth flow is not
triggered. Expand the matcher to also handle direct Snowflake DBAPI exceptions
(not only SQLAlchemy wrappers). [api mismatch]
<details>
<summary><b>Severity Level:</b> Critical π¨</summary>
```mdx
- β SQL Lab queries fail after token revocation.
- β οΈ Dashboards using Snowflake OAuth donβt auto re-auth.
```
</details>
<details>
<summary><b>Steps of Reproduction β
</b></summary>
```mdx
1. Configure Snowflake OAuth2 for Superset so that
`SnowflakeEngineSpec.supports_oauth2`
is True and a client is registered in `DATABASE_OAUTH2_CLIENTS`, and enable
database-level
OAuth (see `Database.is_oauth2_enabled` and `get_oauth2_config` at
`superset/models/core.py:1360-1385`). Run an interactive SQL Lab query
against this
Snowflake database (entry point `superset/sql_lab.py:495-513` where
`database.get_raw_connection` is used).
2. Ensure the user has an OAuth2 access token stored (via the OAuth dance),
but the token
becomes invalid server-side (revoked or otherwise rejected) while still
present in
Superset. The query execution path uses `database.get_raw_connection`
(`superset/models/core.py:619-49`) and then `execute_sql_with_cursor`
(`superset/sql/execution/executor.py:96-79`), which internally calls
`database.db_engine_spec.execute(cursor, stmt_sql, database)` for each
statement.
3. Inside `BaseEngineSpec.execute`
(`superset/db_engine_specs/base.py:2169-27`),
`cursor.execute(query)` is invoked directly on the Snowflake DBAPI cursor.
When the token
is invalid, the Snowflake connector raises a raw
`snowflake.connector.errors.DatabaseError` with a message containing
"Invalid OAuth access
token". The `except Exception as ex` block catches this DBAPI exception and
evaluates
`database.is_oauth2_enabled() and cls.needs_oauth2(ex)`.
4. For Snowflake, `cls.oauth2_exception` is `CustomSnowflakeAuthError`
(`superset/db_engine_specs/snowflake.py:92-93`), whose metaclass
`CustomSnowflakeAuthErrorMeta.__instancecheck__` at
`superset/db_engine_specs/snowflake.py:84-89` returns True only when
`instance` is a
`SqlalchemyDatabaseError` whose `.orig` is `DatabaseError` and whose string
contains
"Invalid OAuth access token`. Since the exception `ex` here is a direct DBAPI
`DatabaseError`, not a `SqlalchemyDatabaseError`, `isinstance(ex,
cls.oauth2_exception)`
is False, so `BaseEngineSpec.needs_oauth2(ex)`
(`superset/db_engine_specs/base.py:2195-10`) returns False. As a result,
OAuth2 re-auth is
not triggered; `BaseEngineSpec.execute` falls through to `raise
cls.get_dbapi_mapped_exception(ex)`, and the user sees a query failure
instead of being
sent through `start_oauth2_dance` to refresh/reauthorize their Snowflake
OAuth token.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=4e7b9437d78c4cb1a912ae12862bd4ef&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=4e7b9437d78c4cb1a912ae12862bd4ef&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:** 85:89
**Comment:**
*Api Mismatch: The OAuth error matcher only recognizes
SQLAlchemy-wrapped `DatabaseError` instances with `.orig`, but query execution
paths use raw DBAPI cursor exceptions as well. When token failures surface as
direct Snowflake connector errors, this check returns false and the OAuth
re-auth flow is not triggered. Expand the matcher to also handle direct
Snowflake DBAPI exceptions (not only SQLAlchemy wrappers).
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=0201d40c7a7fb692920f1e2748854e3a1953c81650fa770e3b15c96ce7999d7c&reaction=like'>π</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=0201d40c7a7fb692920f1e2748854e3a1953c81650fa770e3b15c96ce7999d7c&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]
