diff --git a/app/database/db.py b/app/database/db.py index 0aab09a49..9ad116b17 100644 --- a/app/database/db.py +++ b/app/database/db.py @@ -136,6 +136,15 @@ def _init_file_db(db_url): db_path = get_db_path() db_exists = os.path.exists(db_path) + # Acquire the OS-level lock before any migration work. The lock guards a + # separate `.lock` file, so it cannot block Alembic's own SQLite + # connection. Taking it first is deliberate: revision inspection, backup + # creation, the upgrade and the failure-path restore then run mutually + # exclusively between processes instead of racing each other. + # This diverges from upstream master, which locks after migrating and + # justifies that with a claim about blocking Alembic that is not true. + _acquire_file_lock(db_path) + config = get_alembic_config() # Check if we need to upgrade @@ -177,11 +186,7 @@ def _init_file_db(db_url): logging.exception("Error upgrading database: ") raise e - # Acquire an OS-level file lock after migrations are complete. - # Alembic uses its own connection, so we must wait until it's done - # before locking — otherwise our own lock blocks the migration. conn.close() - _acquire_file_lock(db_path) global Session Session = sessionmaker(bind=engine) diff --git a/tests-unit/app_test/test_db_init_locking.py b/tests-unit/app_test/test_db_init_locking.py new file mode 100644 index 000000000..3594874b9 --- /dev/null +++ b/tests-unit/app_test/test_db_init_locking.py @@ -0,0 +1,71 @@ +"""The DB file lock must cover the migration block, not just follow it. + +A second process arriving while the lock is held has to bail out *before* +inspecting revisions, copying a backup, or running the upgrade — otherwise two +starts against one database race the backup and the restore. +""" +import os +import sqlite3 + +import pytest +from alembic import command +from alembic.config import Config +from filelock import FileLock + +from app.database import db as db_module + +_PRE_HEAD = "0006_add_loader_path" + + +def _make_config(db_path: str) -> Config: + root = os.path.join(os.path.dirname(__file__), "../..") + cfg = Config(os.path.abspath(os.path.join(root, "alembic.ini"))) + cfg.set_main_option("script_location", os.path.abspath(os.path.join(root, "alembic_db"))) + cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}") + return cfg + + +def _current_revision(db_path: str) -> str: + with sqlite3.connect(db_path) as conn: + rows = conn.execute("SELECT version_num FROM alembic_version").fetchall() + assert len(rows) == 1 + return rows[0][0] + + +@pytest.fixture +def stale_db(tmp_path, monkeypatch): + """A file-backed DB parked one revision behind head, wired into args.""" + db_path = str(tmp_path / "comfyui.db") + command.upgrade(_make_config(db_path), _PRE_HEAD) + + monkeypatch.setattr(db_module.args, "database_url", f"sqlite:///{db_path}") + monkeypatch.setattr(db_module, "Session", None) + monkeypatch.setattr(db_module, "_db_lock", None) + yield db_path + if db_module._db_lock is not None: + db_module._db_lock.release(force=True) + + +def test_init_file_db_migrates_when_lock_is_free(stale_db): + """Positive control: unblocked, this same fixture really does migrate.""" + db_module._init_file_db(db_module.args.database_url) + + assert _current_revision(stale_db) != _PRE_HEAD + assert os.path.exists(stale_db + ".bkp") + + +def test_held_lock_blocks_before_any_migration_work(stale_db): + # Given: another process already holds the database's lock file + holder = FileLock(stale_db + ".lock") + holder.acquire(timeout=0) + try: + # When: a second init runs against the same database + with pytest.raises(RuntimeError, match="Another ComfyUI process may already be using it"): + db_module._init_file_db(db_module.args.database_url) + + # Then: it bailed out before backing up or upgrading anything + assert not os.path.exists(stale_db + ".bkp") + assert _current_revision(stale_db) == _PRE_HEAD + assert db_module.Session is None + finally: + holder.release()