fix: update dataset/query catalog on DB changes (#32829)

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,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:
             current_username = get_username()
             SyncPermissionsCommand(
@@ -159,6 +172,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(

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 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()