This is an automated email from the ASF dual-hosted git repository. michaelsmolina pushed a commit to branch 5.0 in repository https://gitbox.apache.org/repos/asf/superset.git
commit c63a33e76b43a3e8ebf28bd80e0cc6372965abad Author: Beto Dealmeida <[email protected]> AuthorDate: Wed Mar 26 08:56:02 2025 -0400 fix: update dataset/query catalog on DB changes (#32829) --- superset/commands/database/update.py | 44 +++++++++++++-- tests/unit_tests/commands/databases/update_test.py | 62 ++++++++++++++++++++++ 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py index fbf90694f4..4a81fba7b5 100644 --- a/superset/commands/database/update.py +++ b/superset/commands/database/update.py @@ -23,7 +23,7 @@ from typing import Any from flask_appbuilder.models.sqla import Model -from superset import is_feature_enabled, security_manager +from superset import db, is_feature_enabled, security_manager from superset.commands.base import BaseCommand from superset.commands.database.exceptions import ( DatabaseConnectionFailedError, @@ -80,13 +80,26 @@ class UpdateDatabaseCommand(BaseCommand): # existing personal tokens. self._handle_oauth2() - # if the database name changed we need to update any existing permissions, - # since they're name based + # build new DB original_database_name = self._model.database_name - + original_catalog = self._model.get_default_catalog() database = DatabaseDAO.update(self._model, self._properties) database.set_sqlalchemy_uri(database.sqlalchemy_uri) ssh_tunnel = self._handle_ssh_tunnel(database) + new_catalog = database.get_default_catalog() + + # update assets when the database catalog changes, if the database was not + # configured with multi-catalog support; if it was enabled or is enabled in the + # update we don't update the assets + if ( + new_catalog != original_catalog + and not self._model.allow_multi_catalog + and not database.allow_multi_catalog + ): + self._update_catalog_attribute(self._model.id, new_catalog) + + # if the database name changed we need to update any existing permissions, + # since they're name based try: self._refresh_catalogs(database, original_database_name, ssh_tunnel) except OAuth2RedirectError: @@ -153,6 +166,29 @@ class UpdateDatabaseCommand(BaseCommand): ssh_tunnel_properties, ).run() + def _update_catalog_attribute( + self, + database_id: int, + new_catalog: str | None, + ) -> None: + """ + Update the catalog of the datasets that are associated with database. + """ + from superset.connectors.sqla.models import SqlaTable + from superset.models.sql_lab import Query, SavedQuery, TableSchema, TabState + + for model in [ + SqlaTable, + Query, + SavedQuery, + TabState, + TableSchema, + ]: + fk = "db_id" if model == SavedQuery else "database_id" + predicate = {fk: database_id} + update = {"catalog": new_catalog} + db.session.query(model).filter_by(**predicate).update(update) + def _get_catalog_names( self, database: Database, diff --git a/tests/unit_tests/commands/databases/update_test.py b/tests/unit_tests/commands/databases/update_test.py index daf41b7506..7c9224bf5a 100644 --- a/tests/unit_tests/commands/databases/update_test.py +++ b/tests/unit_tests/commands/databases/update_test.py @@ -495,3 +495,65 @@ def test_update_other_fields_dont_affect_oauth( add_permission_view_menu.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()
