codeant-ai-for-open-source[bot] commented on code in PR #41421:
URL: https://github.com/apache/superset/pull/41421#discussion_r3476314646
##########
tests/unit_tests/db_engine_specs/test_databricks.py:
##########
@@ -291,3 +300,538 @@ def test_get_prequeries(mocker: MockerFixture) -> None:
"USE CATALOG `evil`` USE CATALOG bad`",
"USE SCHEMA `evil`` USE SCHEMA bad`",
]
+
+
+# OAuth2 Tests
+
+
+def test_oauth2_attributes() -> None:
+ """
+ Test that OAuth2 attributes are properly set for both engine specs.
+ """
+ # Test DatabricksNativeEngineSpec
+ assert DatabricksNativeEngineSpec.supports_oauth2 is True
+ assert DatabricksNativeEngineSpec.oauth2_exception is OAuth2RedirectError
+ assert DatabricksNativeEngineSpec.oauth2_scope == "sql"
+ # OAuth2 endpoints are now dynamic and set at runtime
+ assert DatabricksNativeEngineSpec.oauth2_authorization_request_uri == ""
+ assert DatabricksNativeEngineSpec.oauth2_token_request_uri == ""
+
+ # Test DatabricksPythonConnectorEngineSpec
+ assert DatabricksPythonConnectorEngineSpec.supports_oauth2 is True
+ assert DatabricksPythonConnectorEngineSpec.oauth2_exception is
OAuth2RedirectError
+ assert DatabricksPythonConnectorEngineSpec.oauth2_scope == "sql"
+ # OAuth2 endpoints are now dynamic and set at runtime
+ assert
DatabricksPythonConnectorEngineSpec.oauth2_authorization_request_uri == ""
+ assert DatabricksPythonConnectorEngineSpec.oauth2_token_request_uri == ""
+
+
+def test_impersonate_user_with_token(mocker: MockerFixture) -> None:
+ """
+ Test impersonate_user method with OAuth2 token for
DatabricksNativeEngineSpec.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks+connector://token:original-token@host:443/database"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test with user token
+ url, kwargs = DatabricksNativeEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # Check that the password (token) was updated in the URL
+ assert url.password == "user-oauth-token" # noqa: S105
+ # Check that access_token was updated in connect_args
+ assert kwargs["connect_args"]["access_token"] == "user-oauth-token" #
noqa: S105
+
+
+def test_impersonate_user_without_token(mocker: MockerFixture) -> None:
+ """
+ Test impersonate_user method without OAuth2 token.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+ "databricks+connector://token:original-token@host:443/database"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test without user token
+ url, kwargs = DatabricksNativeEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token=None,
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # Check that nothing was changed
+ assert url.password == "original-token" # noqa: S105
+ assert kwargs["connect_args"]["access_token"] == "original-token" # noqa:
S105
+
+
+def test_impersonate_user_python_connector(mocker: MockerFixture) -> None:
+ """
+ Test impersonate_user method for DatabricksPythonConnectorEngineSpec.
+ """
+ database = mocker.MagicMock()
+ original_url = make_url(
+
"databricks://token:original-token@host:443?http_path=path&catalog=main&schema=default"
+ )
+ engine_kwargs = {"connect_args": {"access_token": "original-token"}}
+
+ # Test with user token
+ url, kwargs = DatabricksPythonConnectorEngineSpec.impersonate_user(
+ database=database,
+ username="user1",
+ user_token="user-oauth-token", # noqa: S106
+ url=original_url,
+ engine_kwargs=engine_kwargs,
+ )
+
+ # Check that the password (token) was updated in the URL
+ assert url.password == "user-oauth-token" # noqa: S105
+ # Check that access_token was updated in connect_args
+ assert kwargs["connect_args"]["access_token"] == "user-oauth-token" #
noqa: S105
+
+
[email protected]
+def oauth2_config_native() -> OAuth2ClientConfig:
+ """
+ Config for Databricks Native OAuth2.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "authorization_request_uri":
"https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/authorize";,
+ "token_request_uri":
"https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token";,
+ "request_content_type": "json",
+ }
+
+
[email protected]
+def oauth2_config_python() -> OAuth2ClientConfig:
+ """
+ Config for Databricks Python Connector OAuth2.
+ """
+ return {
+ "id": "databricks-client-id",
+ "secret": "databricks-client-secret",
+ "scope": "sql",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "authorization_request_uri":
"https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/authorize";,
+ "token_request_uri":
"https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token";,
+ "request_content_type": "json",
+ }
+
+
+def test_is_oauth2_enabled_no_config_native(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is not configured for Native engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={"DATABASE_OAUTH2_CLIENTS": {}},
+ )
+
+ assert DatabricksNativeEngineSpec.is_oauth2_enabled() is False
+
+
+def test_is_oauth2_enabled_config_native(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is configured for Native engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={
+ "DATABASE_OAUTH2_CLIENTS": {
+ "Databricks (legacy)": {
+ "id": "client-id",
+ "secret": "client-secret",
+ },
+ }
+ },
+ )
+
+ assert DatabricksNativeEngineSpec.is_oauth2_enabled() is True
+
+
+def test_is_oauth2_enabled_no_config_python(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is not configured for Python
Connector engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={"DATABASE_OAUTH2_CLIENTS": {}},
+ )
+
+ assert DatabricksPythonConnectorEngineSpec.is_oauth2_enabled() is False
+
+
+def test_is_oauth2_enabled_config_python(mocker: MockerFixture) -> None:
+ """
+ Test `is_oauth2_enabled` when OAuth2 is configured for Python Connector
engine.
+ """
+ mocker.patch(
+ "flask.current_app.config",
+ new={
+ "DATABASE_OAUTH2_CLIENTS": {
+ "Databricks": {
+ "id": "client-id",
+ "secret": "client-secret",
+ },
+ }
+ },
+ )
+
+ assert DatabricksPythonConnectorEngineSpec.is_oauth2_enabled() is True
+
+
+def test_get_oauth2_authorization_uri_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_authorization_uri` for Native engine.
+ """
+ from superset.db_engine_specs.base import OAuth2State
+
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri":
"http://localhost:8088/api/v1/database/oauth2/";,
+ "tab_id": "1234",
+ }
+
+ url = DatabricksNativeEngineSpec.get_oauth2_authorization_uri(
+ oauth2_config_native, state
+ )
+ parsed = urlparse(url)
+ assert parsed.netloc == "accounts.cloud.databricks.com"
+ assert parsed.path == "/oidc/accounts/12345/v1/authorize"
+
+ query = parse_qs(parsed.query)
+ assert query["scope"][0] == "sql"
+ encoded_state = query["state"][0].replace("%2E", ".")
+ assert decode_oauth2_state(encoded_state) == state
+
+
+def test_get_oauth2_authorization_uri_python(
+ mocker: MockerFixture,
+ oauth2_config_python: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_authorization_uri` for Python Connector engine.
+ """
+ from superset.db_engine_specs.base import OAuth2State
+
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri":
"http://localhost:8088/api/v1/database/oauth2/";,
+ "tab_id": "1234",
+ }
+
+ url = DatabricksPythonConnectorEngineSpec.get_oauth2_authorization_uri(
+ oauth2_config_python, state
+ )
+ parsed = urlparse(url)
+ assert parsed.netloc == "accounts.cloud.databricks.com"
+ assert parsed.path == "/oidc/accounts/12345/v1/authorize"
+
+ query = parse_qs(parsed.query)
+ assert query["scope"][0] == "sql"
+ encoded_state = query["state"][0].replace("%2E", ".")
+ assert decode_oauth2_state(encoded_state) == state
+
+
+def test_get_oauth2_token_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_token` for Native engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+
+ assert DatabricksNativeEngineSpec.get_oauth2_token(
+ oauth2_config_native, "authorization-code"
+ ) == {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token";,
+ json={
+ "code": "authorization-code",
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "grant_type": "authorization_code",
+ },
+ timeout=30.0,
+ )
+
+
+def test_get_oauth2_token_python(
+ mocker: MockerFixture,
+ oauth2_config_python: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_token` for Python Connector engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+
+ assert DatabricksPythonConnectorEngineSpec.get_oauth2_token(
+ oauth2_config_python, "authorization-code"
+ ) == {
+ "access_token": "access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token";,
+ json={
+ "code": "authorization-code",
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "redirect_uri": "http://localhost:8088/api/v1/database/oauth2/";,
+ "grant_type": "authorization_code",
+ },
+ timeout=30.0,
+ )
+
+
+def test_get_oauth2_fresh_token_native(
+ mocker: MockerFixture,
+ oauth2_config_native: OAuth2ClientConfig,
+) -> None:
+ """
+ Test `get_oauth2_fresh_token` for Native engine.
+ """
+ requests = mocker.patch("superset.db_engine_specs.base.requests")
+ requests.post().json.return_value = {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+
+ assert DatabricksNativeEngineSpec.get_oauth2_fresh_token(
+ oauth2_config_native, "old-refresh-token"
+ ) == {
+ "access_token": "new-access-token",
+ "expires_in": 3600,
+ "scope": "sql",
+ "token_type": "Bearer",
+ "refresh_token": "new-refresh-token",
+ }
+ requests.post.assert_called_with(
+ "https://accounts.cloud.databricks.com/oidc/accounts/12345/v1/token";,
+ json={
+ "client_id": "databricks-client-id",
+ "client_secret": "databricks-client-secret",
+ "refresh_token": "old-refresh-token",
+ "grant_type": "refresh_token",
+ },
+ timeout=30.0,
+ )
+
+
+def _oauth2_state() -> OAuth2State:
+ state: OAuth2State = {
+ "database_id": 1,
+ "user_id": 1,
+ "default_redirect_uri":
"http://localhost:8088/api/v1/database/oauth2/";,
+ "tab_id": "1234",
+ }
+ return state
Review Comment:
**Suggestion:** Add a concise docstring to this new helper function
describing that it builds and returns the default OAuth2 state used by tests.
[custom_rule]
**Severity Level:** Minor ⚠️
<details>
<summary><b>Why it matters? 🤔 </b></summary>
The helper function was newly added in the final file state and does not
include a docstring. This matches the custom rule requiring newly added Python
functions to be documented inline, so the suggestion is verified.
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=fdbd2a11af1346e0a8afac76e55596d4&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=fdbd2a11af1346e0a8afac76e55596d4&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_databricks.py
**Line:** 668:675
**Comment:**
*Custom Rule: Add a concise docstring to this new helper function
describing that it builds and returns the default OAuth2 state used by tests.
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%2F41421&comment_hash=26de82b294530c84e6481b863d7a2775a62c03a3560e63629c60e14d8b0c4454&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F41421&comment_hash=26de82b294530c84e6481b863d7a2775a62c03a3560e63629c60e14d8b0c4454&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]
