diff --git a/mealie/services/backups_v2/alchemy_exporter.py b/mealie/services/backups_v2/alchemy_exporter.py index 3e4b0877e..95dc1d37a 100644 --- a/mealie/services/backups_v2/alchemy_exporter.py +++ b/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 @@ from mealie.services._base_service import BaseService 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 @@ class AlchemyExporter(BaseService): 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) + 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)); -""" + 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() diff --git a/tests/data/__init__.py b/tests/data/__init__.py index 2b6be2b56..2972f9184 100644 --- a/tests/data/__init__.py +++ b/tests/data/__init__.py @@ -16,15 +16,18 @@ backup_version_44e8d670719d_3 = CWD / "backups/backup-version-44e8d670719d-3.zip 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" diff --git a/tests/data/backups/backup-version-86054b40fd06-1.zip b/tests/data/backups/backup-version-86054b40fd06-1.zip new file mode 100644 index 000000000..baca9748f Binary files /dev/null and b/tests/data/backups/backup-version-86054b40fd06-1.zip differ diff --git a/tests/unit_tests/services_tests/backup_v2_tests/test_backup_v2.py b/tests/unit_tests/services_tests/backup_v2_tests/test_backup_v2.py index 8dac15f4e..cb1af7ae7 100644 --- a/tests/unit_tests/services_tests/backup_v2_tests/test_backup_v2.py +++ b/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):