rusackas commented on code in PR #36856:
URL: https://github.com/apache/superset/pull/36856#discussion_r3565173448


##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +242,105 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
         ),
     }
 
+    # OAuth 2.0 support
+    supports_oauth2 = True

Review Comment:
   Applied — `supports_oauth2: bool = True` (8ba0c75b39).



##########
superset/db_engine_specs/snowflake.py:
##########
@@ -191,6 +242,105 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
         ),
     }
 
+    # OAuth 2.0 support
+    supports_oauth2 = True
+    oauth2_exception = CustomSnowflakeAuthError

Review Comment:
   Applied — `oauth2_exception: type[Exception] = CustomSnowflakeAuthError` 
(8ba0c75b39).



##########
tests/unit_tests/db_engine_specs/test_snowflake.py:
##########
@@ -438,3 +439,177 @@ 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 = mocker.patch("superset.db_engine_specs.base.requests")

Review Comment:
   Applied — `requests: mock.MagicMock = mocker.patch(...)` (8ba0c75b39).



##########
tests/unit_tests/db_engine_specs/test_snowflake.py:
##########
@@ -438,3 +439,177 @@ 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 = 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(mocker: MockerFixture) -> None:
+    """
+    Test that Snowflake supports user impersonation.
+    """
+    from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
+    from superset.models.core import Database
+
+    database = Database(sqlalchemy_uri="snowflake://abc")

Review Comment:
   Applied — `database: Database = Database(...)` (8ba0c75b39).



-- 
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]

Reply via email to