blob: 32b816295be5fc8b6626e8e2d826658334bd6f14 [file] [log] [blame]
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from unittest.mock import MagicMock
from pytest_mock import MockerFixture
from superset import db
from superset.commands.database.update import UpdateDatabaseCommand
from superset.extensions import security_manager
from superset.utils import json
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: MagicMock,
) -> None:
"""
Test that permissions are updated correctly.
In this test, the database has two catalogs with two schemas each:
- catalog1
- schema1
- schema2
- catalog2
- schema3
- schema4
When update is called, only `catalog2.schema3` has permissions associated with it,
so `catalog1.*` and `catalog2.schema4` are added.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None, # first catalog is new
"[my_db].[catalog2]", # second catalog already exists
"[my_db].[catalog2].[schema3]", # first schema already exists
None, # second schema is new
# these are called when checking for existing perms in [db].[schema] format
None,
None,
]
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_pvm.assert_has_calls(
[
# first catalog is added with all schemas
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(
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,
) -> None:
"""
Test that permissions are updated correctly.
In this test, the database has no catalogs and two schemas:
- schema1
- schema2
When update is called, only `schema2` has permissions associated with it, so `schema1`
is added.
""" # noqa: E501
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
f"[{database_without_catalog.name}].[schema1]",
)
def test_rename_with_catalog(
mocker: MockerFixture,
database_with_catalog: MagicMock,
) -> None:
"""
Test that permissions are renamed correctly.
In this test, the database has two catalogs with two schemas each:
- catalog1
- schema1
- schema2
- catalog2
- schema3
- schema4
When update is called, only `catalog2.schema3` has permissions associated with it,
so `catalog1.*` and `catalog2.schema4` are added. Additionally, the database has
been renamed from `my_db` to `my_other_db`.
"""
original_database = mocker.MagicMock()
original_database.database_name = "my_db"
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"
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()
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
"[my_db].[catalog2]", # second catalog already exists
"[my_db].[catalog2].[schema3]", # first schema already exists
None, # second schema is new
# these are called when renaming the permissions:
catalog2_pvm, # old [my_db].[catalog2]
catalog2_schema3_pvm, # old [my_db].[catalog2].[schema3]
None, # [my_db].[catalog2].[schema4] doesn't exist
]
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_pvm.assert_has_calls(
[
# first catalog is added with all schemas with the new DB name
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(
db.session,
security_manager,
"schema_access",
f"[{database_with_catalog.name}].[catalog2].[schema4]",
),
],
)
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 == 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(
mocker: MockerFixture,
database_without_catalog: MockerFixture,
) -> None:
"""
Test that permissions are renamed correctly.
In this test, the database has no catalogs and two schemas:
- schema1
- schema2
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
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
original_database = mocker.MagicMock()
original_database.database_name = "my_db"
database_without_catalog.database_name = "my_other_db"
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_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_pvm.assert_called_with(
db.session,
security_manager,
"schema_access",
f"[{database_without_catalog.name}].[schema1]",
)
assert schema2_pvm.view_menu.name == f"[{database_without_catalog.name}].[schema2]"
def test_rename_without_catalog_with_assets(
mocker: MockerFixture,
database_without_catalog: MockerFixture,
) -> None:
"""
Test that permissions are renamed correctly when the DB connection does not support
catalogs, and it has assets associated with it.
"""
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
original_database = mocker.MagicMock()
original_database.database_name = "my_db"
database_without_catalog.database_name = "my_other_db"
database_without_catalog.get_all_schema_names.return_value = ["schema1"]
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
mocker.patch("superset.commands.database.update.get_username")
mocker.patch("superset.security_manager.get_user_by_username")
dataset = mocker.MagicMock()
chart = mocker.MagicMock()
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",
)
schema_pvm = mocker.MagicMock()
schema_pvm.view_menu.name = "[my_db].[schema1]"
find_permission_view_menu.side_effect = [
"[my_db].[schema1]",
schema_pvm,
]
UpdateDatabaseCommand(1, {}).run()
assert schema_pvm.view_menu.name == f"[{database_without_catalog.name}].[schema1]"
assert dataset.schema_perm == f"[{database_without_catalog.name}].[schema1]"
assert dataset.catalog_perm is None
assert chart.catalog_perm is None
assert chart.schema_perm == f"[{database_without_catalog.name}].[schema1]"
def test_update_with_oauth2(
mocker: MockerFixture,
database_needs_oauth2: MockerFixture,
) -> None:
"""
Test that the database can be updated even if OAuth2 is needed to connect.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {}).run()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_not_called()
def test_update_with_oauth2_changed(
mocker: MockerFixture,
database_needs_oauth2: MockerFixture,
) -> None:
"""
Test that the database can be updated even if OAuth2 is needed to connect.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None, # schema1 has no permissions
"[my_db].[schema2]", # second schema already exists
]
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"
UpdateDatabaseCommand(
1,
{
"masked_encrypted_extra": json.dumps(
{"oauth2_client_info": modified_oauth2_client_info}
)
},
).run()
add_pvm.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.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None,
"[my_db].[schema2]",
]
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {"masked_encrypted_extra": None}).run()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
UpdateDatabaseCommand(1, {"masked_encrypted_extra": "{}"}).run()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
def test_update_oauth2_removes_masked_encrypted_extra_key(
mocker: MockerFixture,
database_needs_oauth2: MockerFixture,
) -> None:
"""
Test that the ``masked_encrypted_extra`` key is properly purged from the properties.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None,
"[my_db].[schema2]",
]
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"
UpdateDatabaseCommand(
1,
{
"masked_encrypted_extra": json.dumps(
{"oauth2_client_info": modified_oauth2_client_info}
)
},
).run()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_called()
database_dao.update.assert_called_with(
database_needs_oauth2,
{
"encrypted_extra": json.dumps(
{"oauth2_client_info": modified_oauth2_client_info}
)
},
)
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.
"""
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,
"find_permission_view_menu",
)
find_permission_view_menu.side_effect = [
None,
"[my_db].[schema2]",
]
add_pvm = mocker.patch("superset.commands.database.sync_permissions.add_pvm")
UpdateDatabaseCommand(1, {"database_name": "New DB name"}).run()
add_pvm.assert_not_called()
database_needs_oauth2.purge_oauth2_tokens.assert_not_called()
def test_update_with_catalog_change(mocker: MockerFixture) -> None:
"""
Test that assets are updated when the main catalog changes.
"""
old_database = mocker.MagicMock(allow_multi_catalog=False)
old_database.get_default_catalog.return_value = "project-A"
old_database.id = 1
new_database = mocker.MagicMock(allow_multi_catalog=False)
new_database.get_default_catalog.return_value = "project-B"
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = old_database
database_dao.update.return_value = new_database
mocker.patch("superset.commands.database.update.SyncPermissionsCommand")
mocker.patch.object(
UpdateDatabaseCommand,
"validate",
)
update_catalog_attribute = mocker.patch.object(
UpdateDatabaseCommand,
"_update_catalog_attribute",
)
UpdateDatabaseCommand(1, {}).run()
update_catalog_attribute.assert_called_once_with(1, "project-B")
def test_update_without_catalog_change(mocker: MockerFixture) -> None:
"""
Test that assets are not updated when the main catalog doesn't change.
"""
old_database = mocker.MagicMock(allow_multi_catalog=False)
old_database.database_name = "Ye Old DB"
old_database.get_default_catalog.return_value = "project-A"
old_database.id = 1
new_database = mocker.MagicMock(allow_multi_catalog=False)
new_database.database_name = "Fancy new DB"
new_database.get_default_catalog.return_value = "project-A"
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = old_database
database_dao.update.return_value = new_database
mocker.patch("superset.commands.database.update.SyncPermissionsCommand")
mocker.patch.object(
UpdateDatabaseCommand,
"validate",
)
update_catalog_attribute = mocker.patch.object(
UpdateDatabaseCommand,
"_update_catalog_attribute",
)
UpdateDatabaseCommand(1, {}).run()
update_catalog_attribute.assert_not_called()
def test_update_broken_connection(mocker: MockerFixture) -> None:
"""
Test that updating a database with a broken connection works
even if it has to run a query to get the default catalog.
"""
database = mocker.MagicMock()
database.get_default_catalog.side_effect = Exception("Broken connection")
database.id = 1
new_db = mocker.MagicMock()
new_db.get_default_catalog.return_value = "main"
database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO")
database_dao.find_by_id.return_value = database
database_dao.update.return_value = new_db
mocker.patch("superset.commands.database.update.SyncPermissionsCommand")
update_catalog_attribute = mocker.patch.object(
UpdateDatabaseCommand,
"_update_catalog_attribute",
)
UpdateDatabaseCommand(1, {}).run()
update_catalog_attribute.assert_called_once_with(1, "main")