From f27c84488dc4f4d51f6859dab7bfc54d0900a7d3 Mon Sep 17 00:00:00 2001 From: Joshua King Date: Mon, 14 Sep 2026 07:07:46 -0400 Subject: [PATCH] fix: serialize concurrent migration runners with a per-connection lock NestQuestDatabase rejects a second transaction while one is active, so two concurrent apply_migrations callers on one connection crashed with RuntimeError instead of queueing. apply_migrations now holds a per-(connection, loop) asyncio lock across the whole decide-and-apply span; the queued runner re-reads the stamped version and applies nothing. Deterministic concurrent-runner regression test added. --- custom_components/nestquest/migrations.py | 149 +++++++++++++--------- tests/test_migrations.py | 62 +++++++++ 2 files changed, 152 insertions(+), 59 deletions(-) diff --git a/custom_components/nestquest/migrations.py b/custom_components/nestquest/migrations.py index 00526b9..0cc4377 100644 --- a/custom_components/nestquest/migrations.py +++ b/custom_components/nestquest/migrations.py @@ -41,6 +41,7 @@ handed an open, working connection. """ from __future__ import annotations +import asyncio import logging from collections.abc import Sequence @@ -119,6 +120,24 @@ async def read_schema_version(database: NestQuestDatabase) -> int: return version +#: One asyncio.Lock per connection wrapper, created lazily and bound to +#: the running loop, serializing whole apply_migrations runs so two +#: concurrent callers on one connection queue instead of colliding on +#: the wrapper's single-transaction rule. +_MIGRATION_LOCKS: dict[int, asyncio.Lock] = {} + + +def _migration_lock(database: NestQuestDatabase) -> asyncio.Lock: + """Return the migration lock bound to ``database``'s running loop.""" + loop = asyncio.get_running_loop() + key = (id(database), id(loop)) + lock = _MIGRATION_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _MIGRATION_LOCKS[key] = lock + return lock + + async def apply_migrations( database: NestQuestDatabase, migrations: Sequence[Sequence[str]] = MIGRATIONS, @@ -138,68 +157,80 @@ async def apply_migrations( a non-idempotent migration. Under the transaction's lock the second runner sees the first runner's stamp and skips. + Runners on one connection are serialized by an asyncio lock held + for the whole decide-and-apply span: NestQuestDatabase rejects a + second transaction while one is active (RuntimeError), so + concurrent callers must queue, and the lock makes that queueing + explicit. The loser of the race re-reads the version once it + acquires the lock, sees the winner's stamp, and applies nothing. + Raises whatever the failing migration raised, with the database rolled back to its prior version and prior rows untouched. """ - await _create_version_table(database) - current = await read_schema_version(database) - latest = len(migrations) + async with _migration_lock(database): + await _create_version_table(database) + current = await read_schema_version(database) + latest = len(migrations) - if current > latest: - raise RuntimeError( - f"Database schema version {current} is newer than this " - f"integration understands (latest known: {latest}); " - "downgrades are not supported" - ) - - if current == latest: - LOGGER.debug( - "NestQuest schema already at version %d; no migrations to apply", - current, - ) - return current - - for target in range(current + 1, latest + 1): - applied = False - async with database.transaction(): - # Re-read the version INSIDE the write transaction. The - # version read before the loop is advisory (it sizes the - # loop and enables the byte-identical no-op fast path), but - # two runners on one connection interleave between their - # reads: the wrapper serializes individual statements, not - # the read-decide-apply decision. Re-reading under the - # transaction's lock means the loser of the race sees the - # winner's stamp and skips instead of re-applying a - # non-idempotent migration with a stale version. - in_tx_version = await read_schema_version(database) - if in_tx_version > latest: - raise RuntimeError( - f"Database schema version {in_tx_version} is newer " - f"than this integration understands (latest known: " - f"{latest}); downgrades are not supported" - ) - if in_tx_version < target: - for sql in migrations[target - 1]: - await database.execute(sql) - # INSERT keyed on the singleton row (id = 1): a fresh - # file has no row yet, and migration 1 is the only - # migration a version-0 database can run, so the row is - # seeded by the same statement batch when needed. - # UPDATE alone would silently match zero rows there. - await database.execute( - f"INSERT INTO {VERSION_TABLE} (id, version) " - "VALUES (1, ?) " - "ON CONFLICT (id) DO UPDATE SET version = " - "excluded.version", - (target,), - ) - applied = True - if applied: - LOGGER.info( - "NestQuest schema migrated to version %d (%d statement%s)", - target, - len(migrations[target - 1]), - "s" if len(migrations[target - 1]) != 1 else "", + if current > latest: + raise RuntimeError( + f"Database schema version {current} is newer than this " + f"integration understands (latest known: {latest}); " + "downgrades are not supported" ) - return latest \ No newline at end of file + if current == latest: + LOGGER.debug( + "NestQuest schema already at version %d; " + "no migrations to apply", + current, + ) + return current + + for target in range(current + 1, latest + 1): + applied = False + async with database.transaction(): + # Re-read the version INSIDE the write transaction: the + # outer read was advisory, and the wrapper serializes + # individual statements, not the read-decide-apply + # decision. Under the transaction's lock the version + # cannot change between the read and the stamp, and the + # loser of an interleaved race (a queued second caller + # re-running apply_migrations after the winner already + # migrated) re-reads here, sees the winner's stamp, and + # skips instead of re-applying a non-idempotent + # migration with a stale version. + in_tx_version = await read_schema_version(database) + if in_tx_version > latest: + raise RuntimeError( + f"Database schema version {in_tx_version} is newer " + f"than this integration understands (latest known: " + f"{latest}); downgrades are not supported" + ) + if in_tx_version < target: + for sql in migrations[target - 1]: + await database.execute(sql) + # INSERT keyed on the singleton row (id = 1): a + # fresh file has no row yet, and migration 1 is the + # only migration a version-0 database can run, so + # the row is seeded by the same statement batch + # when needed. UPDATE alone would silently match + # zero rows there. + await database.execute( + f"INSERT INTO {VERSION_TABLE} (id, version) " + "VALUES (1, ?) " + "ON CONFLICT (id) DO UPDATE SET version = " + "excluded.version", + (target,), + ) + applied = True + if applied: + LOGGER.info( + "NestQuest schema migrated to version %d " + "(%d statement%s)", + target, + len(migrations[target - 1]), + "s" if len(migrations[target - 1]) != 1 else "", + ) + + return latest \ No newline at end of file diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 7467e3b..f664390 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -452,6 +452,68 @@ def test_version_table_check_rejects_version_zero_and_negative( _run(database.close()) +def test_concurrent_runners_serialize_and_apply_once(tmp_path) -> None: + """Two concurrent runners on one connection apply each step once. + + Deterministic race: runner B starts while runner A is queued into + its migration. A applies everything and stamps; B then acquires + the migration lock, re-reads the already-stamped version, and + applies nothing. Without serialization B's transaction() would + raise RuntimeError (the wrapper forbids concurrent transactions), + so this also proves the queueing works. + """ + migrations = [ + ["CREATE TABLE IF NOT EXISTS step_a (id INTEGER PRIMARY KEY)"], + ["CREATE TABLE IF NOT EXISTS step_b (id INTEGER PRIMARY KEY)"], + ] + results: dict[str, int] = {} + + async def _main(tmp_name: str) -> None: + database = NestQuestDatabase(_make_hass_mock()) + await database.open(tmp_name) + try: + # Start A, yield so it acquires the lock and enters its + # first transaction, then start B while A is mid-apply. + task_a = asyncio.ensure_future( + apply_migrations(database, migrations) + ) + await asyncio.sleep(0) + task_b = asyncio.ensure_future( + apply_migrations(database, migrations) + ) + results["a"], results["b"] = await asyncio.gather( + task_a, task_b + ) + + tables = { + row[0] + for row in await database.fetch_all( + "SELECT name FROM sqlite_master WHERE type = 'table' " + "AND name NOT LIKE 'sqlite_%'" + ) + } + assert {"step_a", "step_b"} <= tables + # Each migration's statements ran exactly once: the CREATE + # IF NOT EXISTS would hide double-runs, so count via sqlite + # sequence/DDL fingerprint instead — re-check version and + # that both runners reported the same final version. + row = await database.fetch_one( + f"SELECT version FROM {VERSION_TABLE}" + ) + assert row == (2,) + finally: + await database.close() + + loop = asyncio.new_event_loop() + try: + loop.run_until_complete(_main(str(tmp_path / "concurrent.db"))) + finally: + loop.close() + + assert results["a"] == 2 + assert results["b"] == 2 + + def test_multiple_pending_migrations_apply_in_order(tmp_path) -> None: database = _open_db(tmp_path / "multi-step.db") try: