fix: make catalog migration lenient (#29549)

tests/unit_tests/migrations/shared/catalogs_test.py

@@ -143,3 +143,128 @@ def test_upgrade_catalog_perms(mocker: MockerFixture, session: Session) -> None:
         ("[my_db].[public]",),
         ("[my_db].[db]",),
     ]
+
+
+def test_upgrade_catalog_perms_graceful(
+    mocker: MockerFixture,
+    session: Session,
+) -> None:
+    """
+    Test the `upgrade_catalog_perms` function when it fails to connect to the DB.
+
+    During the migration we try to connect to the analytical database to get the list of
+    schemas. This should fail gracefully and not raise an exception, since the database
+    could be offline, and the permissions can be generated later then the admin enables
+    catalog browsing on the database (permissions are always synced on a DB update, see
+    `UpdateDatabaseCommand`).
+    """
+    from superset.connectors.sqla.models import SqlaTable
+    from superset.models.core import Database
+    from superset.models.slice import Slice
+    from superset.models.sql_lab import Query, SavedQuery, TableSchema, TabState
+
+    engine = session.get_bind()
+    Database.metadata.create_all(engine)
+
+    mocker.patch("superset.migrations.shared.catalogs.op")
+    db = mocker.patch("superset.migrations.shared.catalogs.db")
+    db.Session.return_value = session
+
+    mocker.patch.object(
+        Database,
+        "get_all_schema_names",
+        side_effect=Exception("Failed to connect to the database"),
+    )
+    mocker.patch("superset.migrations.shared.catalogs.op", session)
+
+    database = Database(
+        database_name="my_db",
+        sqlalchemy_uri="postgresql://localhost/db",
+    )
+    dataset = SqlaTable(
+        table_name="my_table",
+        database=database,
+        catalog=None,
+        schema="public",
+        schema_perm="[my_db].[public]",
+    )
+    session.add(dataset)
+    session.commit()
+
+    chart = Slice(
+        slice_name="my_chart",
+        datasource_type="table",
+        datasource_id=dataset.id,
+    )
+    query = Query(
+        client_id="foo",
+        database=database,
+        catalog=None,
+        schema="public",
+    )
+    saved_query = SavedQuery(
+        database=database,
+        sql="SELECT * FROM public.t",
+        catalog=None,
+        schema="public",
+    )
+    tab_state = TabState(
+        database=database,
+        catalog=None,
+        schema="public",
+    )
+    table_schema = TableSchema(
+        database=database,
+        catalog=None,
+        schema="public",
+    )
+    session.add_all([chart, query, saved_query, tab_state, table_schema])
+    session.commit()
+
+    # before migration
+    assert dataset.catalog is None
+    assert query.catalog is None
+    assert saved_query.catalog is None
+    assert tab_state.catalog is None
+    assert table_schema.catalog is None
+    assert dataset.schema_perm == "[my_db].[public]"
+    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]",),
+    ]
+
+    upgrade_catalog_perms()
+
+    # after migration
+    assert dataset.catalog == "db"
+    assert query.catalog == "db"
+    assert saved_query.catalog == "db"
+    assert tab_state.catalog == "db"
+    assert table_schema.catalog == "db"
+    assert dataset.schema_perm == "[my_db].[db].[public]"
+    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]",),
+    ]
+
+    downgrade_catalog_perms()
+
+    # revert
+    assert dataset.catalog is None
+    assert query.catalog is None
+    assert saved_query.catalog is None
+    assert tab_state.catalog is None
+    assert table_schema.catalog is None
+    assert dataset.schema_perm == "[my_db].[public]"
+    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]",),
+    ]