fix: catalog upgrade/downgrade (#29780)

tests/unit_tests/migrations/shared/catalogs_test.py

@@ -22,17 +22,18 @@ from superset.migrations.shared.catalogs import (
     downgrade_catalog_perms,
     upgrade_catalog_perms,
 )
-from superset.migrations.shared.security_converge import ViewMenu
+from superset.migrations.shared.security_converge import (
+    Permission,
+    PermissionView,
+    ViewMenu,
+)
 
 
 def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
     """
     Test the `upgrade_catalog_perms` function.
 
-    The function is called when catalogs are introduced into a new DB engine spec. When
-    that happens, we need to update the `catalog` attribute so it points to the default
-    catalog, instead of being `NULL`. We also need to update `schema_perms` to include
-    the default catalog.
+    The function is called when catalogs are introduced into a new DB engine spec.
     """
     from superset.connectors.sqla.models import SqlaTable
     from superset.models.core import Database
@@ -51,6 +52,11 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
         "get_all_schema_names",
         return_value=["public", "information_schema"],
     )
+    mocker.patch.object(
+        Database,
+        "get_all_catalog_names",
+        return_value=["db", "other_catalog"],
+    )
 
     database = Database(
         database_name="my_db",
@@ -61,6 +67,7 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
         database=database,
         catalog=None,
         schema="public",
+        catalog_perm=None,
         schema_perm="[my_db].[public]",
     )
     session.add(dataset)
@@ -70,6 +77,8 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
         slice_name="my_chart",
         datasource_type="table",
         datasource_id=dataset.id,
+        catalog_perm=None,
+        schema_perm="[my_db].[public]",
     )
     query = Query(
         client_id="foo",
@@ -102,15 +111,43 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
     assert saved_query.catalog is None
     assert tab_state.catalog is None
     assert table_schema.catalog is None
+    assert dataset.catalog_perm is None
     assert dataset.schema_perm == "[my_db].[public]"
+    assert chart.catalog_perm is None
     assert chart.schema_perm == "[my_db].[public]"
-    assert session.query(ViewMenu.name).all() == [
-        ("[my_db].(id:1)",),
-        ("[my_db].[my_table](id:1)",),
-        ("[my_db].[public]",),
+    assert (
+        session.query(ViewMenu.name, Permission.name)
+        .join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
+        .join(Permission, PermissionView.permission_id == Permission.id)
+        .all()
+    ) == [
+        ("[my_db].(id:1)", "database_access"),
+        ("[my_db].[my_table](id:1)", "datasource_access"),
+        ("[my_db].[public]", "schema_access"),
     ]
 
     upgrade_catalog_perms()
+    session.commit()
+
+    # add dataset/chart in new catalog
+    new_dataset = SqlaTable(
+        table_name="my_table",
+        database=database,
+        catalog="other_catalog",
+        schema="public",
+        schema_perm="[my_db].[other_catalog].[public]",
+        catalog_perm="[my_db].[other_catalog]",
+    )
+    session.add(new_dataset)
+    session.commit()
+
+    new_chart = Slice(
+        slice_name="my_chart",
+        datasource_type="table",
+        datasource_id=new_dataset.id,
+    )
+    session.add(new_chart)
+    session.commit()
 
     # after migration
     assert dataset.catalog == "db"
@@ -118,16 +155,29 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
     assert saved_query.catalog == "db"
     assert tab_state.catalog == "db"
     assert table_schema.catalog == "db"
+    assert dataset.catalog_perm == "[my_db].[db]"
     assert dataset.schema_perm == "[my_db].[db].[public]"
+    assert chart.catalog_perm == "[my_db].[db]"
     assert chart.schema_perm == "[my_db].[db].[public]"
-    assert session.query(ViewMenu.name).all() == [
-        ("[my_db].(id:1)",),
-        ("[my_db].[my_table](id:1)",),
-        ("[my_db].[db].[public]",),
-        ("[my_db].[db]",),
+    assert (
+        session.query(ViewMenu.name, Permission.name)
+        .join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
+        .join(Permission, PermissionView.permission_id == Permission.id)
+        .all()
+    ) == [
+        ("[my_db].(id:1)", "database_access"),
+        ("[my_db].[my_table](id:1)", "datasource_access"),
+        ("[my_db].[db].[public]", "schema_access"),
+        ("[my_db].[db]", "catalog_access"),
+        ("[my_db].[other_catalog]", "catalog_access"),
+        ("[my_db].[other_catalog].[public]", "schema_access"),
+        ("[my_db].[other_catalog].[information_schema]", "schema_access"),
+        ("[my_db].[my_table](id:2)", "datasource_access"),
     ]
 
+    # do a downgrade
     downgrade_catalog_perms()
+    session.commit()
 
     # revert
     assert dataset.catalog is None
@@ -135,15 +185,25 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
     assert saved_query.catalog is None
     assert tab_state.catalog is None
     assert table_schema.catalog is None
+    assert dataset.catalog_perm is None
     assert dataset.schema_perm == "[my_db].[public]"
+    assert chart.catalog_perm is None
     assert chart.schema_perm == "[my_db].[public]"
-    assert session.query(ViewMenu.name).all() == [
-        ("[my_db].(id:1)",),
-        ("[my_db].[my_table](id:1)",),
-        ("[my_db].[public]",),
-        ("[my_db].[db]",),
+    assert (
+        session.query(ViewMenu.name, Permission.name)
+        .join(PermissionView, ViewMenu.id == PermissionView.view_menu_id)
+        .join(Permission, PermissionView.permission_id == Permission.id)
+        .all()
+    ) == [
+        ("[my_db].(id:1)", "database_access"),
+        ("[my_db].[my_table](id:1)", "datasource_access"),
+        ("[my_db].[public]", "schema_access"),
     ]
 
+    # make sure new dataset/chart were deleted
+    assert session.query(SqlaTable).all() == [dataset]
+    assert session.query(Slice).all() == [chart]
+
 
 def test_upgrade_catalog_perms_graceful(
     mocker: MockerFixture,
@@ -236,6 +296,7 @@ def test_upgrade_catalog_perms_graceful(
     ]
 
     upgrade_catalog_perms()
+    session.commit()
 
     # after migration
     assert dataset.catalog == "db"
@@ -253,6 +314,7 @@ def test_upgrade_catalog_perms_graceful(
     ]
 
     downgrade_catalog_perms()
+    session.commit()
 
     # revert
     assert dataset.catalog is None
@@ -266,5 +328,4 @@ def test_upgrade_catalog_perms_graceful(
         ("[my_db].(id:1)",),
         ("[my_db].[my_table](id:1)",),
         ("[my_db].[public]",),
-        ("[my_db].[db]",),
     ]