This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch update-catalog-on-change in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6483f88b896ca63ac6d89f65ccf3c9cd84ce6f6b Author: Beto Dealmeida <[email protected]> AuthorDate: Mon Mar 24 15:18:51 2025 -0400 fix: update dataset/query catalog on DB changes --- superset/commands/database/update.py | 50 +++++++++++++++-- tests/unit_tests/commands/databases/update_test.py | 62 ++++++++++++++++++++++ 2 files changed, 107 insertions(+), 5 deletions(-) diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py index a1accf4df1..2dfdffd68b 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 +from superset import db, is_feature_enabled from superset.commands.base import BaseCommand from superset.commands.database.exceptions import ( DatabaseExistsValidationError, @@ -79,13 +79,30 @@ 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 - original_database_name = self._model.database_name - + # build new DB + old_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 + print("BETO HERE", old_catalog, new_catalog) + if ( + not self._model.allow_multi_catalog + and not database.allow_multi_catalog + and new_catalog != old_catalog + ): + self._update_catalog_attribute( + self._model.id, + database.get_default_catalog(), + ) + + # if the database name changed we need to update any existing permissions, + # since they're name based + original_database_name = self._model.database_name try: current_username = get_username() SyncPermissionsCommand( @@ -159,6 +176,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 validate(self) -> None: if database_name := self._properties.get("database_name"): if not DatabaseDAO.validate_update_uniqueness( diff --git a/tests/unit_tests/commands/databases/update_test.py b/tests/unit_tests/commands/databases/update_test.py index a74ff3c5c0..4c70234774 100644 --- a/tests/unit_tests/commands/databases/update_test.py +++ b/tests/unit_tests/commands/databases/update_test.py @@ -580,3 +580,65 @@ def test_update_other_fields_dont_affect_oauth( 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 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()
