QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-18 07:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Update database permissions in async mode (#32231)

Browse Source
This commit is contained in:
Vitor Avila
2025-02-28 21:25:47 -03:00
committed by GitHub
parent 84b52b2323
commit d79f7b28c2
22 changed files with 1715 additions and 425 deletions
Download Patch File Download Diff File Expand all files Collapse all files

367
tests/unit_tests/commands/databases/update_test.py
View File

@@ -17,84 +17,19 @@
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from superset import db
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",
"secret": "client_secret",
"scope": "scope-a",
"redirect_uri": "redirect_uri",
"authorization_request_uri": "auth_uri",
"token_request_uri": "token_uri",
"request_content_type": "json",
}
@pytest.fixture
def database_with_catalog(mocker: MockerFixture) -> MagicMock:
"""
Mock a database with catalogs and schemas.
"""
database = mocker.MagicMock()
database.database_name = "my_db"
database.db_engine_spec.__name__ = "test_engine"
database.db_engine_spec.supports_catalog = True
database.get_all_catalog_names.return_value = ["catalog1", "catalog2"]
database.get_all_schema_names.side_effect = [
["schema1", "schema2"],
["schema3", "schema4"],
]
database.get_default_catalog.return_value = "catalog2"
return database
@pytest.fixture
def database_without_catalog(mocker: MockerFixture) -> MagicMock:
"""
Mock a database without catalogs.
"""
database = mocker.MagicMock()
database.database_name = "my_db"
database.db_engine_spec.__name__ = "test_engine"
database.db_engine_spec.supports_catalog = False
database.get_all_schema_names.return_value = ["schema1", "schema2"]
return database
@pytest.fixture
def database_needs_oauth2(mocker: MockerFixture) -> MagicMock:
"""
Mock a database without catalogs that needs OAuth2.
"""
database = mocker.MagicMock()
database.database_name = "my_db"
database.db_engine_spec.__name__ = "test_engine"
database.db_engine_spec.supports_catalog = False
database.get_all_schema_names.side_effect = OAuth2RedirectError(
"url",
"tab_id",
"redirect_uri",
)
database.encrypted_extra = json.dumps({"oauth2_client_info": oauth2_client_info})
database.db_engine_spec.unmask_encrypted_extra = (
BaseEngineSpec.unmask_encrypted_extra
)
return database
from tests.conftest import with_config
from tests.unit_tests.commands.databases.conftest import oauth2_client_info
def test_update_with_catalog(
mocker: MockerFixture,
database_with_catalog: MockerFixture,
database_with_catalog: MagicMock,
) -> None:
"""
Test that permissions are updated correctly.
@@ -111,9 +46,15 @@ def test_update_with_catalog(
When update is called, only `catalog2.schema3` has permissions associated with it,
so `catalog1.*` and `catalog2.schema4` are added.
"""
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
DatabaseDAO.find_by_id.return_value = database_with_catalog
DatabaseDAO.update.return_value = database_with_catalog
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_with_catalog
database_dao.update.return_value = database_with_catalog
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_with_catalog
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -128,25 +69,66 @@ def test_update_with_catalog(
None,
None,
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_has_calls(
add_pvm.assert_has_calls(
[
# first catalog is added with all schemas
mocker.call("catalog_access", "[my_db].[catalog1]"),
mocker.call("schema_access", "[my_db].[catalog1].[schema1]"),
mocker.call("schema_access", "[my_db].[catalog1].[schema2]"),
mocker.call(
db.session, security_manager, "catalog_access", "[my_db].[catalog1]"
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_db].[catalog1].[schema1]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_db].[catalog1].[schema2]",
),
# second catalog already exists, only `schema4` is added
mocker.call("schema_access", "[my_db].[catalog2].[schema4]"),
mocker.call(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog2].[schema4]",
),
],
)
@with_config({"SYNC_DB_PERMISSIONS_IN_ASYNC_MODE": True})
def test_update_sync_perms_in_async_mode(
mocker: MockerFixture,
database_with_catalog: MagicMock,
) -> None:
"""
Test that updating a DB connection with async mode enables
triggers the celery task to syn perms.
"""
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_with_catalog
database_dao.update.return_value = database_with_catalog
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_with_catalog
sync_task = mocker.patch(
"superset.commands.database.sync_permissions.sync_database_permissions_task.delay"
)
mocker.patch("superset.commands.database.update.get_username", return_value="admin")
mocker.patch("superset.security_manager.get_user_by_username")
UpdateDatabaseCommand(1, {}).run()
sync_task.assert_called_once_with(1, "admin", "my_db")
def test_update_without_catalog(
mocker: MockerFixture,
database_without_catalog: MockerFixture,
@@ -162,9 +144,15 @@ def test_update_without_catalog(
When update is called, only `schema2` has permissions associated with it, so `schema1`
is added.
""" # noqa: E501
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
DatabaseDAO.find_by_id.return_value = database_without_catalog
DatabaseDAO.update.return_value = database_without_catalog
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_without_catalog
database_dao.update.return_value = database_without_catalog
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_without_catalog
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -174,22 +162,21 @@ def test_update_without_catalog(
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_called_with(
add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
"[my_db].[schema1]",
f"[{database_without_catalog.name}].[schema1]",
)
def test_rename_with_catalog(
mocker: MockerFixture,
database_with_catalog: MockerFixture,
database_with_catalog: MagicMock,
) -> None:
"""
Test that permissions are renamed correctly.
@@ -207,25 +194,33 @@ def test_rename_with_catalog(
so `catalog1.*` and `catalog2.schema4` are added. Additionally, the database has
been renamed from `my_db` to `my_other_db`.
"""
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
original_database = mocker.MagicMock()
original_database.database_name = "my_db"
DatabaseDAO.find_by_id.return_value = original_database
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = original_database
database_with_catalog.database_name = "my_other_db"
DatabaseDAO.update.return_value = database_with_catalog
database_dao.update.return_value = database_with_catalog
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_with_catalog
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
dataset = mocker.MagicMock()
chart = mocker.MagicMock()
DatabaseDAO.get_datasets.return_value = [dataset]
DatasetDAO = mocker.patch("superset.commands.database.update.DatasetDAO") # noqa: N806
DatasetDAO.get_related_objects.return_value = {"charts": [chart]}
sync_db_perms_dao.get_datasets.return_value = [dataset]
dataset_dao = mocker.patch("superset.commands.database.sync_permissions.DatasetDAO")
dataset_dao.get_related_objects.return_value = {"charts": [chart]}
find_permission_view_menu = mocker.patch.object(
security_manager,
"find_permission_view_menu",
)
catalog2_pvm = mocker.MagicMock()
catalog2_pvm.view_menu.name = "[my_db].[catalog2]"
catalog2_schema3_pvm = mocker.MagicMock()
catalog2_schema3_pvm.view_menu.name = "[my_db].[catalog2].[schema3]"
find_permission_view_menu.side_effect = [
# these are called when adding the permissions:
None, # first catalog is new
@@ -237,31 +232,52 @@ def test_rename_with_catalog(
catalog2_schema3_pvm, # old [my_db].[catalog2].[schema3]
None, # [my_db].[catalog2].[schema4] doesn't exist
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
add_vm = mocker.patch("superset.commands.database.sync_permissions.add_vm")
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_has_calls(
add_pvm.assert_has_calls(
[
# first catalog is added with all schemas with the new DB name
mocker.call("catalog_access", "[my_other_db].[catalog1]"),
mocker.call("schema_access", "[my_other_db].[catalog1].[schema1]"),
mocker.call("schema_access", "[my_other_db].[catalog1].[schema2]"),
mocker.call(
db.session,
security_manager,
"catalog_access",
"[my_other_db].[catalog1]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_other_db].[catalog1].[schema1]",
),
mocker.call(
db.session,
security_manager,
"schema_access",
"[my_other_db].[catalog1].[schema2]",
),
# second catalog already exists, only `schema4` is added
mocker.call("schema_access", "[my_other_db].[catalog2].[schema4]"),
mocker.call(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog2].[schema4]",
),
],
)
assert catalog2_pvm.view_menu.name == "[my_other_db].[catalog2]"
assert catalog2_schema3_pvm.view_menu.name == "[my_other_db].[catalog2].[schema3]"
assert catalog2_pvm.view_menu == add_vm.return_value
assert (
catalog2_schema3_pvm.view_menu.name
== f"[{database_with_catalog.name}].[catalog2].[schema3]"
)
assert dataset.catalog_perm == "[my_other_db].[catalog2]"
assert dataset.schema_perm == "[my_other_db].[catalog2].[schema4]"
assert chart.catalog_perm == "[my_other_db].[catalog2]"
assert chart.schema_perm == "[my_other_db].[catalog2].[schema4]"
assert dataset.catalog_perm == f"[{database_with_catalog.name}].[catalog2]"
assert dataset.schema_perm == f"[{database_with_catalog.name}].[catalog2].[schema4]"
assert chart.catalog_perm == f"[{database_with_catalog.name}].[catalog2]"
assert chart.schema_perm == f"[{database_with_catalog.name}].[catalog2].[schema4]"
def test_rename_without_catalog(
@@ -279,38 +295,44 @@ def test_rename_without_catalog(
When update is called, only `schema2` has permissions associated with it, so `schema1`
is added. Additionally, the database has been renamed from `my_db` to `my_other_db`.
""" # noqa: E501
DatabaseDAO = mocker.patch("superset.commands.database.update.DatabaseDAO") # noqa: N806
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
original_database = mocker.MagicMock()
original_database.database_name = "my_db"
DatabaseDAO.find_by_id.return_value = original_database
database_without_catalog.database_name = "my_other_db"
DatabaseDAO.update.return_value = database_without_catalog
DatabaseDAO.get_datasets.return_value = []
database_dao.update.return_value = database_without_catalog
database_dao.find_by_id.return_value = original_database
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_without_catalog
sync_db_perms_dao.get_datasets.return_value = []
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
"find_permission_view_menu",
)
schema2_pvm = mocker.MagicMock()
schema2_pvm.view_menu.name = "[my_db].[schema2]"
find_permission_view_menu.side_effect = [
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
None, # [my_db].[schema1] doesn't exist
schema2_pvm, # old [my_db].[schema2]
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_called_with(
add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
"[my_other_db].[schema1]",
f"[{database_without_catalog.name}].[schema1]",
)
assert schema2_pvm.view_menu.name == "[my_other_db].[schema2]"
assert schema2_pvm.view_menu.name == f"[{database_without_catalog.name}].[schema2]"
def test_update_with_oauth2(
@@ -320,9 +342,15 @@ def test_update_with_oauth2(
"""
Test that the database can be updated even if OAuth2 is needed to connect.
"""
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_needs_oauth2
database_dao.update.return_value = database_needs_oauth2
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_needs_oauth2
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -332,14 +360,11 @@ def test_update_with_oauth2(
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_not_called()
@@ -350,9 +375,15 @@ def test_update_with_oauth2_changed(
"""
Test that the database can be updated even if OAuth2 is needed to connect.
"""
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_needs_oauth2
database_dao.update.return_value = database_needs_oauth2
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_needs_oauth2
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -362,10 +393,7 @@ def test_update_with_oauth2_changed(
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
modified_oauth2_client_info = oauth2_client_info.copy()
modified_oauth2_client_info["scope"] = "scope-b"
@@ -379,7 +407,7 @@ def test_update_with_oauth2_changed(
},
).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
@@ -390,9 +418,15 @@ def test_remove_oauth_config_purges_tokens(
"""
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_needs_oauth2
database_dao.update.return_value = database_needs_oauth2
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_needs_oauth2
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -402,19 +436,16 @@ def test_remove_oauth_config_purges_tokens(
None,
"[my_db].[schema2]",
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {"masked_encrypted_extra": None}).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
UpdateDatabaseCommand(1, {"masked_encrypted_extra": "{}"}).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
@@ -425,9 +456,15 @@ def test_update_oauth2_removes_masked_encrypted_extra_key(
"""
Test that the ``masked_encrypted_extra`` key is properly purged from the properties.
"""
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_needs_oauth2
database_dao.update.return_value = database_needs_oauth2
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_needs_oauth2
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -437,10 +474,7 @@ def test_update_oauth2_removes_masked_encrypted_extra_key(
None,
"[my_db].[schema2]",
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
modified_oauth2_client_info = oauth2_client_info.copy()
modified_oauth2_client_info["scope"] = "scope-b"
@@ -454,9 +488,9 @@ def test_update_oauth2_removes_masked_encrypted_extra_key(
},
).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
DatabaseDAO.update.assert_called_with(
database_dao.update.assert_called_with(
database_needs_oauth2,
{
"encrypted_extra": json.dumps(
@@ -474,9 +508,15 @@ def test_update_other_fields_dont_affect_oauth(
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database_needs_oauth2
database_dao.update.return_value = database_needs_oauth2
sync_db_perms_dao = mocker.patch(
"superset.commands.database.sync_permissions.DatabaseDAO"
)
sync_db_perms_dao.find_by_id.return_value = database_needs_oauth2
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
find_permission_view_menu = mocker.patch.object(
security_manager,
@@ -486,12 +526,9 @@ def test_update_other_fields_dont_affect_oauth(
None,
"[my_db].[schema2]",
]
add_permission_view_menu = mocker.patch.object(
security_manager,
"add_permission_view_menu",
)
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {"database_name": "New DB name"}).run()
add_permission_view_menu.assert_not_called()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_not_called()