Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 4 commits into from
Jan 10, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 18 additions & 11 deletions superset/commands/database/update.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
from superset.db_engine_specs.base import GenericDBException
from superset.exceptions import OAuth2RedirectError
from superset.models.core import Database
from superset.utils import json
from superset.utils.decorators import on_error, transaction

logger = logging.getLogger(__name__)
Expand All @@ -66,22 +67,23 @@ def run(self) -> Model:

self.validate()

# unmask ``encrypted_extra``
self._properties["encrypted_extra"] = (
self._model.db_engine_spec.unmask_encrypted_extra(
self._model.encrypted_extra,
self._properties.pop("masked_encrypted_extra", "{}"),
if "masked_encrypted_extra" in self._properties:
# unmask ``encrypted_extra``
self._properties["encrypted_extra"] = (
self._model.db_engine_spec.unmask_encrypted_extra(
self._model.encrypted_extra,
self._properties["masked_encrypted_extra"],
)
)
)

# Depending on the changes to the OAuth2 configuration we may need to purge
# existing personal tokens.
self._handle_oauth2()

# if the database name changed we need to update any existing permissions,
# since they're name based
original_database_name = self._model.database_name

# Depending on the changes to the OAuth2 configuration we may need to purge
# existing personal tokens.
self._handle_oauth2()

database = DatabaseDAO.update(self._model, self._properties)
database.set_sqlalchemy_uri(database.sqlalchemy_uri)
ssh_tunnel = self._handle_ssh_tunnel(database)
Expand All @@ -99,11 +101,16 @@ def _handle_oauth2(self) -> None:
if not self._model:
return

if self._properties["encrypted_extra"] is None:
self._model.purge_oauth2_tokens()
return

current_config = self._model.get_oauth2_config()
if not current_config:
return

new_config = self._properties["encrypted_extra"].get("oauth2_client_info", {})
encrypted_extra = json.loads(self._properties["encrypted_extra"])
new_config = encrypted_extra.get("oauth2_client_info", {})

# Keys that require purging personal tokens because they probably are no longer
# valid. For example, if the scope has changed the existing tokens are still
Expand Down
89 changes: 80 additions & 9 deletions tests/unit_tests/commands/databases/update_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,10 @@
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",
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Loading