fix(db): repair model_relationships foreign keys broken by migration 22

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -34,6 +34,7 @@
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_29 import build_migration_29
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_30 import build_migration_30
 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_31 import build_migration_31
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_32 import build_migration_32
 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
 
 
@@ -85,6 +86,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto
     migrator.register_migration(build_migration_29())
     migrator.register_migration(build_migration_30())
     migrator.register_migration(build_migration_31())
+    migrator.register_migration(build_migration_32())
     migrator.run_migrations()
 
     return db

invokeai/app/services/shared/sqlite_migrator/migrations/migration_32.py

@@ -0,0 +1,91 @@
+"""Migration 32: Repair model_relationships foreign keys.
+
+Migration 22 rebuilt the `models` table by renaming it to `models_old`, creating a
+fresh `models` table, copying the data over, and dropping `models_old`. Because
+modern SQLite (with `legacy_alter_table` off) rewrites foreign-key references in
+*other* tables when a table is renamed, the foreign keys in `model_relationships`
+were silently repointed at `models_old` -- which was then dropped.
+
+This left the related-models links referencing a table that no longer exists,
+breaking `ON DELETE CASCADE` and foreign-key integrity for related models.
+
+This migration rebuilds `model_relationships` so its foreign keys reference
+`models(id)` again, preserving existing links and dropping any orphaned rows whose
+model keys no longer exist (those would violate the restored foreign keys).
+"""
+
+import sqlite3
+
+from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
+
+
+class Migration32Callback:
+    """Migration to repair the broken foreign keys on the model_relationships table."""
+
+    def __call__(self, cursor: sqlite3.Cursor) -> None:
+        self._repair_model_relationships_fks(cursor)
+
+    def _repair_model_relationships_fks(self, cursor: sqlite3.Cursor) -> None:
+        cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='model_relationships';")
+        row = cursor.fetchone()
+        if row is None:
+            # Table does not exist (fresh db will create it correctly), nothing to repair.
+            return
+
+        existing_sql: str = row[0]
+        if "models_old" not in existing_sql:
+            # Foreign keys already point at the correct table, nothing to repair.
+            return
+
+        # Rebuild the table with the correct foreign keys referencing models(id).
+        cursor.execute("ALTER TABLE model_relationships RENAME TO model_relationships_old;")
+        cursor.execute(
+            """
+            -- many-to-many relationship table for models
+            CREATE TABLE model_relationships (
+                -- model_key_1 and model_key_2 are the same as the key(primary key) in the models table
+                model_key_1 TEXT NOT NULL,
+                model_key_2 TEXT NOT NULL,
+                created_at TEXT DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
+                PRIMARY KEY (model_key_1, model_key_2),
+                -- model_key_1 < model_key_2, to ensure uniqueness and prevent duplicates
+                FOREIGN KEY (model_key_1) REFERENCES models(id) ON DELETE CASCADE,
+                FOREIGN KEY (model_key_2) REFERENCES models(id) ON DELETE CASCADE
+            );
+            """
+        )
+
+        # Copy over the existing links, dropping any orphaned rows whose model keys no
+        # longer exist -- these would violate the restored foreign keys.
+        cursor.execute(
+            """
+            INSERT INTO model_relationships (model_key_1, model_key_2, created_at)
+            SELECT model_key_1, model_key_2, created_at
+            FROM model_relationships_old
+            WHERE model_key_1 IN (SELECT id FROM models)
+              AND model_key_2 IN (SELECT id FROM models);
+            """
+        )
+
+        # Drop the old table first so its index name is freed before we recreate it.
+        cursor.execute("DROP TABLE model_relationships_old;")
+        cursor.execute(
+            """
+            -- Creates an index to keep performance equal when searching for model_key_1 or model_key_2
+            CREATE INDEX IF NOT EXISTS keyx_model_relationships_model_key_2
+            ON model_relationships(model_key_2);
+            """
+        )
+
+
+def build_migration_32() -> Migration:
+    """Builds the migration object for migrating from version 31 to version 32.
+
+    This migration repairs the foreign keys on the model_relationships table, which were
+    broken by migration 22 rebuilding the models table.
+    """
+    return Migration(
+        from_version=31,
+        to_version=32,
+        callback=Migration32Callback(),
+    )

tests/app/services/shared/sqlite_migrator/migrations/test_migration_32.py

@@ -0,0 +1,131 @@
+"""Tests for migration 32: Repair model_relationships foreign keys."""
+
+import sqlite3
+
+import pytest
+
+from invokeai.app.services.shared.sqlite_migrator.migrations.migration_32 import (
+    Migration32Callback,
+    build_migration_32,
+)
+
+
+def _create_models_table(conn: sqlite3.Connection) -> None:
+    conn.execute(
+        """
+        CREATE TABLE models (
+            id TEXT NOT NULL PRIMARY KEY,
+            config TEXT NOT NULL
+        );
+        """
+    )
+
+
+def _create_broken_relationships_table(conn: sqlite3.Connection) -> None:
+    """Recreates the broken state left by migration 22: FKs reference the dropped models_old table."""
+    conn.execute(
+        """
+        CREATE TABLE model_relationships (
+            model_key_1 TEXT NOT NULL,
+            model_key_2 TEXT NOT NULL,
+            created_at TEXT DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
+            PRIMARY KEY (model_key_1, model_key_2),
+            FOREIGN KEY (model_key_1) REFERENCES "models_old"(id) ON DELETE CASCADE,
+            FOREIGN KEY (model_key_2) REFERENCES "models_old"(id) ON DELETE CASCADE
+        );
+        """
+    )
+    conn.execute("CREATE INDEX keyx_model_relationships_model_key_2 ON model_relationships(model_key_2);")
+
+
+@pytest.fixture
+def db() -> sqlite3.Connection:
+    conn = sqlite3.connect(":memory:")
+    _create_models_table(conn)
+    _create_broken_relationships_table(conn)
+    conn.execute("INSERT INTO models (id, config) VALUES ('a', '{}'), ('b', '{}'), ('c', '{}')")
+    return conn
+
+
+class TestMigration32:
+    def test_repoints_foreign_keys_to_models(self, db: sqlite3.Connection):
+        """After migration, the foreign keys reference models, not models_old."""
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        sql = db.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='model_relationships'").fetchone()[
+            0
+        ]
+        assert "models_old" not in sql
+        assert "REFERENCES models(id)" in sql
+
+    def test_preserves_valid_links(self, db: sqlite3.Connection):
+        """Links between existing models are preserved."""
+        db.execute("INSERT INTO model_relationships (model_key_1, model_key_2) VALUES ('a', 'b')")
+        db.commit()
+
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        rows = db.execute("SELECT model_key_1, model_key_2 FROM model_relationships ORDER BY model_key_1").fetchall()
+        assert rows == [("a", "b")]
+
+    def test_drops_orphaned_links(self, db: sqlite3.Connection):
+        """Links referencing missing models are dropped so the restored FKs are satisfiable."""
+        db.execute("INSERT INTO model_relationships (model_key_1, model_key_2) VALUES ('a', 'b')")
+        db.execute("INSERT INTO model_relationships (model_key_1, model_key_2) VALUES ('a', 'gone')")
+        db.commit()
+
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        rows = db.execute("SELECT model_key_1, model_key_2 FROM model_relationships").fetchall()
+        assert rows == [("a", "b")]
+
+    def test_cascade_works_after_repair(self, db: sqlite3.Connection):
+        """ON DELETE CASCADE against models works once the FKs are repaired."""
+        db.execute("INSERT INTO model_relationships (model_key_1, model_key_2) VALUES ('a', 'b')")
+        db.commit()
+
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        db.execute("PRAGMA foreign_keys = ON;")
+        db.execute("DELETE FROM models WHERE id = 'a'")
+        db.commit()
+
+        rows = db.execute("SELECT * FROM model_relationships").fetchall()
+        assert rows == []
+
+    def test_index_recreated(self, db: sqlite3.Connection):
+        """The lookup index on model_key_2 is recreated on the rebuilt table."""
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        idx = db.execute(
+            "SELECT name FROM sqlite_master WHERE type='index' AND name='keyx_model_relationships_model_key_2'"
+        ).fetchone()
+        assert idx is not None
+
+    def test_idempotent_when_already_correct(self, db: sqlite3.Connection):
+        """Running on an already-correct table is a no-op (no rebuild)."""
+        Migration32Callback()(db.cursor())
+        db.commit()
+        # Second run should detect the correct FKs and do nothing.
+        Migration32Callback()(db.cursor())
+        db.commit()
+
+        sql = db.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='model_relationships'").fetchone()[
+            0
+        ]
+        assert "REFERENCES models(id)" in sql
+
+    def test_no_relationships_table_is_noop(self):
+        """If the table doesn't exist, migration is a no-op."""
+        conn = sqlite3.connect(":memory:")
+        Migration32Callback()(conn.cursor())  # should not raise
+
+    def test_build_migration_32_version_numbers(self):
+        migration = build_migration_32()
+        assert migration.from_version == 31
+        assert migration.to_version == 32