codeant-ai-for-open-source[bot] commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3565175351
##########
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 = instance
Review Comment:
**Suggestion:** Add an explicit type annotation for this local variable to
keep the new metaclass logic fully typed. [custom_rule]
**Severity Level:** Minor π§Ή
<details>
<summary><b>Why it matters? β </b></summary>
The new local variable `orig` can be annotated here, but it is introduced
without a type hint inside the newly added metaclass method. This matches the
Python type-hint rule violation.
</details>
<details>
<summary><b>Rule source π </b></summary>
.cursor/rules/dev-standard.mdc (line 28)
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=82e597dbeddb4f93bbdc9796e6e24020&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=82e597dbeddb4f93bbdc9796e6e24020&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:** 91:91
**Comment:**
*Custom Rule: Add an explicit type annotation for this local variable
to keep the new metaclass logic fully typed.
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=bb19cef628c6ec65976b806a328c04815eb285a3c8d72ca4e13a71517d165dc0&reaction=like'>π</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=bb19cef628c6ec65976b806a328c04815eb285a3c8d72ca4e13a71517d165dc0&reaction=dislike'>π</a>
##########
tests/unit_tests/db_engine_specs/test_snowflake.py:
##########
@@ -438,3 +440,206 @@ def test_unmask_encrypted_extra() -> None:
},
}
)
+
+
[email protected]
+def oauth2_config() -> OAuth2ClientConfig:
+ """
+ Config for Snowflake OAuth2.
+ """
+ return {
+ "id": "snowflake-oauth2-client-id",
+ "secret": "snowflake-oauth2-client-secret",
+ "scope": "refresh_token",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "authorization_request_uri":
"https://snowflake.oauth2.example/oauth/authorize";,
+ "token_request_uri":
"https://snowflake.oauth2.example/oauth/token-request";,
+ "request_content_type": "data",
+ }
+
+
+def test_get_oauth2_token(
+ mocker: MockerFixture,
+ oauth2_config: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_token`.
+ """
+ from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+
+ requests: mock.MagicMock =
mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "scope",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+
+ assert SnowflakeEngineSpec.get_oauth2_token(oauth2_config, "code") == {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "scope",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://snowflake.oauth2.example/oauth/token-request";,
+ data={
+ "code": "code",
+ "client_id": "snowflake-oauth2-client-id",
+ "client_secret": "snowflake-oauth2-client-secret",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "grant_type": "authorization_code",
+ },
+ timeout=30.0,
+ )
+
+
+def test_impersonate_user(app: SupersetApp, mocker: MockerFixture) -> None:
+ """
+ Test that Snowflake supports user impersonation.
+
+ Impersonation only applies within a request context (see
+ ``test_impersonate_user_outside_request_context`` below for the
+ background-execution case), so these assertions run inside one.
+ """
+ from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+ from superset.models.core import Database
+
+ database: Database = Database(sqlalchemy_uri="snowflake://abc")
+
+ mocker.patch(
+
"superset.db_engine_specs.snowflake.SnowflakeEngineSpec.is_oauth2_enabled",
+ return_value=True,
+ )
+
+ with app.test_request_context("/some/place/"):
+ assert SnowflakeEngineSpec.impersonate_user(
+ database=database,
+ username=None,
+ user_token=None,
+
url=make_url("snowflake://user:pass@account/database_name/default"),
+ engine_kwargs={
+ "connect_args": {
+ "validate_default_parameters": True,
+ },
+ },
+ ) == (
+ make_url("snowflake://user:pass@account/database_name/default"),
+ {"connect_args": {"validate_default_parameters": True}},
+ )
+
+ assert SnowflakeEngineSpec.impersonate_user(
+ database=database,
+ username=None,
+ user_token=None,
+
url=make_url("snowflake://user:pass@account/database_name/default"),
+ engine_kwargs={},
+ ) == (
+ make_url(
+
"snowflake://user:pass@account/database_name/default?authenticator=oauth"
+ ),
+ {"connect_args": {"authenticator": "oauth"}},
+ )
+
+ mocker.patch(
+ "superset.db_engine_specs.snowflake.is_feature_enabled",
+ return_value=True,
+ )
+
+ mocker.patch(
+ "superset.security_manager.find_user",
+
return_value=mocker.MagicMock(email="[email protected]"),
+ )
+ assert SnowflakeEngineSpec.impersonate_user(
+ database=database,
+ username="impersonated_user",
+ user_token="test_token", # noqa: S106
+
url=make_url("snowflake://user:pass@account/database_name/default"),
+ engine_kwargs={},
+ ) == (
+ make_url(
+
"snowflake://impersonated_user:pass@account/database_name/default?authenticator=oauth&token=test_token"
+ ),
+ {"connect_args": {"authenticator": "oauth"}},
+ )
+
+
+def test_impersonate_user_outside_request_context(mocker: MockerFixture) ->
None:
+ """
+ Background executions (alerts/reports) have no per-user token, so OAuth
+ impersonation must not engage outside a request context β even when
+ ``database.is_oauth2_enabled()`` returns True because of a
+ database-level OAuth2 client config, which (unlike the app-config-based
+ check) isn't itself request-context-aware.
+ """
+ from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+ from superset.models.core import Database
+
+ database: Database = Database(sqlalchemy_uri="snowflake://abc")
+ mocker.patch.object(Database, "is_oauth2_enabled", return_value=True)
+
+ url = make_url("snowflake://user:pass@account/database_name/default")
Review Comment:
**Suggestion:** Add an explicit type annotation for this local URL variable
to satisfy the requirement for annotating relevant variables in new Python
code. [custom_rule]
**Severity Level:** Minor π§Ή
<details>
<summary><b>Why it matters? β </b></summary>
This is a newly added Python local variable with an obvious concrete type,
but it is not annotated. That matches the rule requiring type hints for
relevant variables in new or modified Python code.
</details>
<details>
<summary><b>Rule source π </b></summary>
.cursor/rules/dev-standard.mdc (line 28)
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=05576e2f9ec2443388ec2c65cb717ada&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=05576e2f9ec2443388ec2c65cb717ada&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:** tests/unit_tests/db_engine_specs/test_snowflake.py
**Line:** 583:583
**Comment:**
*Custom Rule: Add an explicit type annotation for this local URL
variable to satisfy the requirement for annotating relevant variables in new
Python code.
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=591960b5d8bd6b78e5e27b75f9f1687570135a7a1d3b72b379270a5fe775cf74&reaction=like'>π</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=591960b5d8bd6b78e5e27b75f9f1687570135a7a1d3b72b379270a5fe775cf74&reaction=dislike'>π</a>
##########
superset/db_engine_specs/snowflake.py:
##########
@@ -438,6 +596,18 @@ def update_params_from_encrypted_extra(
database: "Database",
params: dict[str, Any],
) -> None:
+ # To use OAuth authentication, a database connection must first be
created using
+ # another authenticator (typically key-pair authentication)
+ # with βImpersonate logged in userβ enabled.
+ # Key-pair authentication is used for connection tests,
+ # while OAuth authentication is used when executing actual queries,
+ # such as in SQL Lab or dashboards.
+ # Therefore, when using OAuth authentication, the key-pair
authentication
+ # settings are not loaded, and the connection is established using
OAuth only.
+ connect_args = params.get("connect_args") or {}
Review Comment:
**Suggestion:** Add a concrete type annotation for this newly added
`connect_args` variable to satisfy the type-hint requirement. [custom_rule]
**Severity Level:** Minor π§Ή
<details>
<summary><b>Why it matters? β </b></summary>
This newly added local variable is assignable to a concrete mapping type and
is introduced without an explicit annotation in modified Python code, so it
violates the type-hint requirement.
</details>
<details>
<summary><b>Rule source π </b></summary>
.cursor/rules/dev-standard.mdc (line 28)
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=3ce1ee70e33140508ba9fdd7d2bad72f&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=3ce1ee70e33140508ba9fdd7d2bad72f&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:** 607:607
**Comment:**
*Custom Rule: Add a concrete type annotation for this newly added
`connect_args` variable to satisfy the type-hint requirement.
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=4eaf5d14d77fc705906e0e2e2619631dd58f0247c76951892e33a6640b0744ad&reaction=like'>π</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F36856&comment_hash=4eaf5d14d77fc705906e0e2e2619631dd58f0247c76951892e33a6640b0744ad&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]
