codeant-ai-for-open-source[bot] commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3486194308
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +224,99 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
),
}
+ # OAuth 2.0 support
+ supports_oauth2 = True
+ oauth2_exception = CustomSnowflakeAuthError
+
+ @classmethod
+ def is_oauth2_enabled(cls) -> bool:
+ """
+ Return whether OAuth2 authentication is enabled.
+ """
+
+ # When alerts or reports connect to the database in the background,
+ # OAuth2 authentication fails; therefore, OAuth2 authentication is
disabled
+ # for background execution.
+ if not has_request_context():
+ return False
+
+ return (
+ cls.supports_oauth2
+ and cls.engine_name in app.config["DATABASE_OAUTH2_CLIENTS"]
+ )
+
+ @classmethod
+ def get_oauth2_config(cls) -> OAuth2ClientConfig | None:
+ """
+ Build the DB engine spec level OAuth2 client config.
+ """
+ if not cls.is_oauth2_enabled():
+ return None
+
+ return super().get_oauth2_config()
+
+ @classmethod
+ def impersonate_user(
+ cls,
+ database: Database,
+ username: str | None,
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ """
+ Modify URL and/or engine kwargs to impersonate a different user.
+ """
+ connect_args = engine_kwargs.setdefault("connect_args", {})
+
+ # When test_connection is executed (i.e., when
validate_default_parameters is
+ # set to True in connect_args), authentication via OAuth is not
performed.
+ if (
+ not connect_args.get("validate_default_parameters", False)
+ and cls.is_oauth2_enabled()
+ ):
+ url = url.update_query_dict({"authenticator": "oauth"})
+ connect_args["authenticator"] = "oauth"
+
+ if user_token and cls.is_oauth2_enabled():
Review Comment:
**Suggestion:** The OAuth gating in impersonation checks only
`cls.is_oauth2_enabled()`, which reads global `DATABASE_OAUTH2_CLIENTS` config
and ignores database-level `oauth2_client_info` stored in `encrypted_extra`. As
a result, Snowflake connections configured with per-database OAuth2 clients can
successfully fetch a user token but never switch the connection to OAuth
authenticator, so queries keep using the non-OAuth auth path. Use the passed
`database` object (for example `database.is_oauth2_enabled()` /
`database.get_oauth2_config()`) for these checks instead of class-level global
config. [api mismatch]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ Snowflake OAuth2 impersonation never activates for database-level
clients.
- ⚠️ User queries run under service identity, breaking impersonation.
- ⚠️ Security audit logs miss correct per-user Snowflake identity.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Configure a Snowflake database object with per-database OAuth2 client
info by setting
`encrypted_extra` to include `oauth2_client_info` so that
`Database.get_oauth2_config()`
in `superset/models/core.py:1333-1352` returns a config, and omit any
`Snowflake` entry
from `app.config["DATABASE_OAUTH2_CLIENTS"]` so that only database-level
configuration is
present.
2. Ensure this Snowflake database has `impersonate_user=True` (column
defined at
`superset/models/core.py:212`) and run a query (for example via SQL Lab),
which routes
through `Database._get_sqla_engine` at `superset/models/core.py:534-589`,
where
`oauth2_config = self.get_oauth2_config()` (line 567) is non-None and
`get_oauth2_access_token(...)` (lines 569-575) retrieves a user access token
from
`DatabaseUserOAuth2Tokens`.
3. In `_get_sqla_engine`, because `self.impersonate_user` is True (line
581), the code
calls `self.db_engine_spec.impersonate_user(self, effective_username,
access_token,
sqlalchemy_url, engine_kwargs)` at `superset/models/core.py:581-588`; for
Snowflake this
resolves to `SnowflakeEngineSpec.impersonate_user` implemented in
`superset/db_engine_specs/snowflake.py:259-292`.
4. Inside `SnowflakeEngineSpec.impersonate_user` (snowflake.py:271-281) both
the branch
that sets `authenticator="oauth"` and the branch that injects the
`user_token` into the
URL are guarded by `cls.is_oauth2_enabled()`, which in
`SnowflakeEngineSpec.is_oauth2_enabled`
(`superset/db_engine_specs/snowflake.py:231-246`)
only checks `cls.engine_name in app.config["DATABASE_OAUTH2_CLIENTS"]` and
ignores the
database’s `oauth2_client_info`; with no global Snowflake entry this returns
False, so the
conditions at lines 274-281 are never entered and the connection continues
to use the
non-OAuth authenticator path even though a per-database OAuth2 client and
user token are
configured.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=3592a871e4e647839dcf38729bdf45e4&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=3592a871e4e647839dcf38729bdf45e4&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:** 274:281
**Comment:**
*Api Mismatch: The OAuth gating in impersonation checks only
`cls.is_oauth2_enabled()`, which reads global `DATABASE_OAUTH2_CLIENTS` config
and ignores database-level `oauth2_client_info` stored in `encrypted_extra`. As
a result, Snowflake connections configured with per-database OAuth2 clients can
successfully fetch a user token but never switch the connection to OAuth
authenticator, so queries keep using the non-OAuth auth path. Use the passed
`database` object (for example `database.is_oauth2_enabled()` /
`database.get_oauth2_config()`) for these checks instead of class-level global
config.
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=3325b90023db062cf366463c78eb67abe45df4cecf95d59e97b48b7c0e9b4e71&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=3325b90023db062cf366463c78eb67abe45df4cecf95d59e97b48b7c0e9b4e71&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]
