fix: Disable Foreign Key Checks During Restore

mealie/services/backups_v2/alchemy_exporter.py

@@ -1,13 +1,15 @@
 import datetime
 import os
 import uuid
+from logging import Logger
 from os import path
 from pathlib import Path
+from textwrap import dedent
 from typing import Any
 
 from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel
-from sqlalchemy import ForeignKey, ForeignKeyConstraint, MetaData, Table, create_engine, insert, text
+from sqlalchemy import Connection, ForeignKey, ForeignKeyConstraint, MetaData, Table, create_engine, insert, text
 from sqlalchemy.engine import base
 from sqlalchemy.orm import sessionmaker
 
@@ -21,6 +23,36 @@
 PROJECT_DIR = Path(__file__).parent.parent.parent.parent
 
 
+class ForeignKeyDisabler:
+    def __init__(self, connection: Connection, dialect_name: str, *, logger: Logger | None = None):
+        self.connection = connection
+        self.is_postgres = dialect_name == "postgresql"
+        self.logger = logger
+
+        self._initial_fk_state: str | None = None
+
+    def __enter__(self):
+        if self.is_postgres:
+            self._initial_fk_state = self.connection.execute(text("SHOW session_replication_role;")).scalar()
+            self.connection.execute(text("SET session_replication_role = 'replica';"))
+        else:
+            self._initial_fk_state = self.connection.execute(text("PRAGMA foreign_keys;")).scalar()
+            self.connection.execute(text("PRAGMA foreign_keys = OFF;"))
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        try:
+            if self.is_postgres:
+                initial_state = self._initial_fk_state or "origin"
+                self.connection.execute(text(f"SET session_replication_role = '{initial_state}';"))
+            else:
+                initial_state = self._initial_fk_state or "ON"
+                self.connection.execute(text(f"PRAGMA foreign_keys = {initial_state};"))
+        except Exception:
+            if self.logger:
+                self.logger.exception("Error when re-enabling foreign keys")
+            raise
+
+
 class AlchemyExporter(BaseService):
     connection_str: str
     engine: base.Engine
@@ -175,40 +207,42 @@ def restore(self, db_dump: dict) -> None:
         del db_dump["alembic_version"]
         """Restores all data from dictionary into the database"""
         with self.engine.begin() as connection:
-            data = self.convert_types(db_dump)
+            with ForeignKeyDisabler(connection, self.engine.dialect.name, logger=self.logger):
+                data = self.convert_types(db_dump)
 
-            self.meta.reflect(bind=self.engine)
-            for table_name, rows in data.items():
-                if not rows:
-                    continue
-                table = self.meta.tables[table_name]
-                rows = self.clean_rows(db_dump, table, rows)
-
-                connection.execute(table.delete())
-                connection.execute(insert(table), rows)
-            if self.engine.dialect.name == "postgresql":
-                # Restore postgres sequence numbers
-                connection.execute(
-                    text(
-                        """
-                SELECT SETVAL('api_extras_id_seq', (SELECT MAX(id) FROM api_extras));
-SELECT SETVAL('group_meal_plans_id_seq', (SELECT MAX(id) FROM group_meal_plans));
-SELECT SETVAL('ingredient_food_extras_id_seq', (SELECT MAX(id) FROM ingredient_food_extras));
-SELECT SETVAL('invite_tokens_id_seq', (SELECT MAX(id) FROM invite_tokens));
-SELECT SETVAL('long_live_tokens_id_seq', (SELECT MAX(id) FROM long_live_tokens));
-SELECT SETVAL('notes_id_seq', (SELECT MAX(id) FROM notes));
-SELECT SETVAL('password_reset_tokens_id_seq', (SELECT MAX(id) FROM password_reset_tokens));
-SELECT SETVAL('recipe_assets_id_seq', (SELECT MAX(id) FROM recipe_assets));
-SELECT SETVAL('recipe_ingredient_ref_link_id_seq', (SELECT MAX(id) FROM recipe_ingredient_ref_link));
-SELECT SETVAL('recipe_nutrition_id_seq', (SELECT MAX(id) FROM recipe_nutrition));
-SELECT SETVAL('recipe_settings_id_seq', (SELECT MAX(id) FROM recipe_settings));
-SELECT SETVAL('recipes_ingredients_id_seq', (SELECT MAX(id) FROM recipes_ingredients));
-SELECT SETVAL('server_tasks_id_seq', (SELECT MAX(id) FROM server_tasks));
-SELECT SETVAL('shopping_list_extras_id_seq', (SELECT MAX(id) FROM shopping_list_extras));
-SELECT SETVAL('shopping_list_item_extras_id_seq', (SELECT MAX(id) FROM shopping_list_item_extras));
-"""
+                self.meta.reflect(bind=self.engine)
+                for table_name, rows in data.items():
+                    if not rows:
+                        continue
+                    table = self.meta.tables[table_name]
+                    rows = self.clean_rows(db_dump, table, rows)
+
+                    connection.execute(table.delete())
+                    connection.execute(insert(table), rows)
+                if self.engine.dialect.name == "postgresql":
+                    # Restore postgres sequence numbers
+                    sequences = [
+                        ("api_extras_id_seq", "api_extras"),
+                        ("group_meal_plans_id_seq", "group_meal_plans"),
+                        ("ingredient_food_extras_id_seq", "ingredient_food_extras"),
+                        ("invite_tokens_id_seq", "invite_tokens"),
+                        ("long_live_tokens_id_seq", "long_live_tokens"),
+                        ("notes_id_seq", "notes"),
+                        ("password_reset_tokens_id_seq", "password_reset_tokens"),
+                        ("recipe_assets_id_seq", "recipe_assets"),
+                        ("recipe_ingredient_ref_link_id_seq", "recipe_ingredient_ref_link"),
+                        ("recipe_nutrition_id_seq", "recipe_nutrition"),
+                        ("recipe_settings_id_seq", "recipe_settings"),
+                        ("recipes_ingredients_id_seq", "recipes_ingredients"),
+                        ("server_tasks_id_seq", "server_tasks"),
+                        ("shopping_list_extras_id_seq", "shopping_list_extras"),
+                        ("shopping_list_item_extras_id_seq", "shopping_list_item_extras"),
+                    ]
+
+                    sql = "\n".join(
+                        [f"SELECT SETVAL('{seq}', (SELECT MAX(id) FROM {table}));" for seq, table in sequences]
                     )
-                )
+                    connection.execute(text(dedent(sql)))
 
         # Re-init database to finish migrations
         init_db.main()

tests/data/__init__.py

@@ -16,15 +16,18 @@
 backup_version_44e8d670719d_4 = CWD / "backups/backup-version-44e8d670719d-4.zip"
 """44e8d670719d: add extras to shopping lists, list items, and ingredient foods"""
 
-backup_version_ba1e4a6cfe99_1 = CWD / "backups/backup-version-ba1e4a6cfe99-1.zip"
-"""ba1e4a6cfe99: added plural names and alias tables for foods and units"""
-
 backup_version_bcfdad6b7355_1 = CWD / "backups/backup-version-bcfdad6b7355-1.zip"
 """bcfdad6b7355: remove tool name and slug unique contraints"""
 
+backup_version_ba1e4a6cfe99_1 = CWD / "backups/backup-version-ba1e4a6cfe99-1.zip"
+"""ba1e4a6cfe99: added plural names and alias tables for foods and units"""
+
 backup_version_09aba125b57a_1 = CWD / "backups/backup-version-09aba125b57a-1.zip"
 """09aba125b57a: add OIDC auth method (Safari-mangled ZIP structure)"""
 
+backup_version_86054b40fd06_1 = CWD / "backups/backup-version-86054b40fd06-1.zip"
+"""86054b40fd06: added query_filter_string to cookbook and mealplan"""
+
 migrations_paprika = CWD / "migrations/paprika.zip"
 
 migrations_chowdown = CWD / "migrations/chowdown.zip"

tests/unit_tests/services_tests/backup_v2_tests/test_backup_v2.py

@@ -84,15 +84,17 @@ def test_database_restore():
         test_data.backup_version_ba1e4a6cfe99_1,
         test_data.backup_version_bcfdad6b7355_1,
         test_data.backup_version_09aba125b57a_1,
+        test_data.backup_version_86054b40fd06_1,
     ],
     ids=[
         "44e8d670719d_1: add extras to shopping lists, list items, and ingredient foods",
         "44e8d670719d_2: add extras to shopping lists, list items, and ingredient foods",
         "44e8d670719d_3: add extras to shopping lists, list items, and ingredient foods",
         "44e8d670719d_4: add extras to shopping lists, list items, and ingredient foods",
-        "ba1e4a6cfe99_1: added plural names and alias tables for foods and units",
         "bcfdad6b7355_1: remove tool name and slug unique contraints",
-        "09aba125b57a: add OIDC auth method (Safari-mangled ZIP structure)",
+        "ba1e4a6cfe99_1: added plural names and alias tables for foods and units",
+        "09aba125b57a_1: add OIDC auth method (Safari-mangled ZIP structure)",
+        "86054b40fd06_1: added query_filter_string to cookbook and mealplan",
     ],
 )
 def test_database_restore_data(backup_path: Path):