fix(oauth): Handle updates to the OAuth config (#31777)

This commit is contained in:
Vitor Avila
2025-01-10 15:54:53 -03:00
committed by GitHub
parent bbdc195a3b
commit 41ed37ab02
2 changed files with 98 additions and 20 deletions

View File

@@ -21,8 +21,10 @@ import pytest
from pytest_mock import MockerFixture
from superset.commands.database.update import UpdateDatabaseCommand
from superset.db_engine_specs.base import BaseEngineSpec
from superset.exceptions import OAuth2RedirectError
from superset.extensions import security_manager
from superset.utils import json
oauth2_client_info = {
"id": "client_id",
@@ -82,7 +84,10 @@ def database_needs_oauth2(mocker: MockerFixture) -> MagicMock:
"tab_id",
"redirect_uri",
)
database.get_oauth2_config.return_value = oauth2_client_info
database.encrypted_extra = json.dumps({"oauth2_client_info": oauth2_client_info})
database.db_engine_spec.unmask_encrypted_extra = (
BaseEngineSpec.unmask_encrypted_extra
)
return database
@@ -332,10 +337,6 @@ def test_update_with_oauth2(
"add_permission_view_menu",
)
database_needs_oauth2.db_engine_spec.unmask_encrypted_extra.return_value = {
"oauth2_client_info": oauth2_client_info,
}
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_not_called()
@@ -368,11 +369,81 @@ def test_update_with_oauth2_changed(
modified_oauth2_client_info = oauth2_client_info.copy()
modified_oauth2_client_info["scope"] = "scope-b"
database_needs_oauth2.db_engine_spec.unmask_encrypted_extra.return_value = {
"oauth2_client_info": modified_oauth2_client_info,
}
UpdateDatabaseCommand(1, {}).run()
UpdateDatabaseCommand(
1,
{
"masked_encrypted_extra": json.dumps(
{"oauth2_client_info": modified_oauth2_client_info}
)
},
).run()
add_permission_view_menu.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
def test_remove_oauth_config_purges_tokens(
mocker: MockerFixture,
database_needs_oauth2: MockerFixture,
) -> None:
"""
Test that removing the OAuth config from a database purges existing tokens.
"""
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
DatabaseDAO.find_by_id.return_value = database_needs_oauth2
DatabaseDAO.update.return_value = database_needs_oauth2
find_permission_view_menu = mocker.patch.object(
security_manager,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None,
"[my_db].[schema2]",
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
UpdateDatabaseCommand(1, {"masked_encrypted_extra": None}).run()
add_permission_view_menu.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
UpdateDatabaseCommand(1, {"masked_encrypted_extra": "{}"}).run()
add_permission_view_menu.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
def test_update_other_fields_dont_affect_oauth(
mocker: MockerFixture,
database_needs_oauth2: MockerFixture,
) -> None:
"""
Test that not including ``masked_encrypted_extra`` in the payload does not
touch the OAuth config.
"""
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
DatabaseDAO.find_by_id.return_value = database_needs_oauth2
DatabaseDAO.update.return_value = database_needs_oauth2
find_permission_view_menu = mocker.patch.object(
security_manager,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None,
"[my_db].[schema2]",
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
UpdateDatabaseCommand(1, {"database_name": "New DB name"}).run()
add_permission_view_menu.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_not_called()