# 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")