This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 68a81c3989 fix: update dataset/query catalog on DB changes (#32829)
68a81c3989 is described below
commit 68a81c39891ea871ea8553f0fc054de0db5d270d
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 a1accf4df1..f562b5d780 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,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(
diff --git a/tests/unit_tests/commands/databases/update_test.py
b/tests/unit_tests/commands/databases/update_test.py
index a74ff3c5c0..5e6b0869d7 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 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()
