fix: Disable Foreign Key Checks During Restore (#4444)

Co-authored-by: Kuchenpirat <24235032+Kuchenpirat@users.noreply.github.com>
This commit is contained in:
Michael Genson 2024-10-29 07:43:57 -05:00 committed by GitHub
parent 3bf6840cbc
commit 8d1ce5c190
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 76 additions and 37 deletions

View file

@ -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()

View file

@ -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"

Binary file not shown.

View file

@ -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):