From 863a144e934e2e051e7ac08edcf17f1193285ea0 Mon Sep 17 00:00:00 2001 From: Joshua King Date: Sun, 13 Sep 2026 21:24:20 -0400 Subject: [PATCH] fix: make transaction lifecycle cancellation-safe and preserve commit+rollback errors --- custom_components/nestquest/db.py | 128 ++++++++++++++++++++---- tests/test_db.py | 157 ++++++++++++++++++++++++++++++ 2 files changed, 267 insertions(+), 18 deletions(-) diff --git a/custom_components/nestquest/db.py b/custom_components/nestquest/db.py index 65163e8..cd1c93b 100644 --- a/custom_components/nestquest/db.py +++ b/custom_components/nestquest/db.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio import sqlite3 -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Awaitable, Sequence from contextlib import asynccontextmanager from dataclasses import dataclass from pathlib import Path @@ -203,6 +203,14 @@ class NestQuestDatabase: failing COMMIT (e.g. a deferred foreign-key violation) also triggers a ROLLBACK before the commit error propagates, so the real connection is never left inside a stray transaction. + + The lifecycle is cancellation-safe. Transaction ownership and the + wrapper lock are only released after the executor work has fully + settled, even when the owning task is cancelled mid-flight: if a + cancelled BEGIN or COMMIT still completes on the worker thread, + the corresponding ROLLBACK is shielded and awaited before any state + is cleared. A lifecycle statement that fails plus a failing + follow-up cleanup surface both errors in an ExceptionGroup. """ if self._in_transaction: raise RuntimeError( @@ -218,7 +226,13 @@ class NestQuestDatabase: def _begin() -> None: conn.execute("BEGIN") - await self._hass.async_add_executor_job(_begin) + await self._settle_lifecycle( + conn, + self._hass.async_add_executor_job(_begin), + "BEGIN did not settle cleanly and the follow-up " + "cleanup also raised; the transaction may need manual " + "cleanup", + ) try: yield except BaseException: @@ -226,31 +240,109 @@ class NestQuestDatabase: def _rollback() -> None: conn.execute("ROLLBACK") - await self._hass.async_add_executor_job(_rollback) + await self._settle_lifecycle( + conn, + self._hass.async_add_executor_job(_rollback), + "ROLLBACK failed and the follow-up cleanup also " + "raised; the transaction may need manual cleanup", + ) raise def _commit() -> None: conn.execute("COMMIT") - try: - await self._hass.async_add_executor_job(_commit) - except BaseException as commit_error: - - def _rollback() -> None: - conn.execute("ROLLBACK") - - try: - await self._hass.async_add_executor_job(_rollback) - except BaseException as rollback_error: - raise RuntimeError( - "COMMIT failed and the follow-up ROLLBACK also " - "raised; the transaction may need manual cleanup" - ) from rollback_error - raise + await self._settle_lifecycle( + conn, + self._hass.async_add_executor_job(_commit), + "COMMIT failed and the follow-up cleanup also raised; " + "the transaction may need manual cleanup", + ) finally: self._in_transaction = False self._tx_task = None + async def _settle_lifecycle( + self, conn: sqlite3.Connection, job: Awaitable[Any], message: str + ) -> None: + """Run one BEGIN/ROLLBACK/COMMIT job, then force the connection clean. + + The job always settles first (see :meth:`_await_settled`), so a + cancelled statement that the worker thread completed anyway is + never abandoned mid-flight. If the job raised, a follow-up + ROLLBACK — itself settled — guarantees the connection ends outside + any SQLite transaction before the error propagates. When both the + job and the cleanup raise, both errors are preserved in an + ExceptionGroup instead of the original being lost. + """ + try: + await self._await_settled(job) + except BaseException as settle_error: + errors: list[BaseException] = [settle_error] + try: + await self._ensure_no_stray_transaction(conn) + except BaseException as cleanup_error: + errors.append(cleanup_error) + if len(errors) == 2: + raise BaseExceptionGroup(message, errors) from settle_error + raise + + async def _ensure_no_stray_transaction( + self, conn: sqlite3.Connection + ) -> None: + """Force the connection out of any transaction via ROLLBACK. + + Runs after a failed or cancelled lifecycle statement. The worker + thread may still be finishing (or may already have completed) the + very statement whose await was cancelled, so both the check and + the ROLLBACK settle in the executor before the caller releases + transaction ownership. + """ + + def _cleanup() -> bool: + if conn.in_transaction: + conn.execute("ROLLBACK") + return conn.in_transaction + + stray = await self._await_settled( + self._hass.async_add_executor_job(_cleanup) + ) + if stray: + raise RuntimeError( + "The SQLite connection is still inside a transaction after " + "cleanup; the transaction may need manual cleanup" + ) + + async def _await_settled(self, job: Awaitable[Any]) -> Any: + """Await ``job``, letting it settle even if this task is cancelled. + + Cancellation is observed, not swallowed: it is re-raised only once + the job has fully settled, so the caller can finalize against a + known connection state instead of abandoning in-flight work. The + job is wrapped in a task because a bare coroutine dies un-awaited + once its await is cancelled. A second cancellation during + settlement is absorbed the same way; if the job itself fails after + a cancellation, both errors surface in an ExceptionGroup. + """ + cancelled = False + inner = asyncio.ensure_future(job) + while not inner.done(): + try: + await asyncio.shield(inner) + except asyncio.CancelledError: + cancelled = True + try: + result = inner.result() + except BaseException as job_error: + if cancelled: + raise BaseExceptionGroup( + "executor job failed while settling after cancellation", + [asyncio.CancelledError(), job_error], + ) from job_error + raise + if cancelled: + raise asyncio.CancelledError() + return result + async def close(self) -> None: """Close the connection in the executor; safe to call repeatedly. diff --git a/tests/test_db.py b/tests/test_db.py index fbd75a1..61d4977 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -406,6 +406,163 @@ async def test_commit_failure_rolls_back_and_next_transaction_works( await database.close() +async def test_cancel_during_begin_settles_then_rolls_back(tmp_path) -> None: + """NQ-PR12-007: cancel during BEGIN must not leak a stray transaction. + + The BEGIN executor job is gated: the transaction task is cancelled + while BEGIN is still in flight, then the worker is allowed to + complete it anyway. The wrapper must keep transaction ownership + until the submitted work settles, roll the completed BEGIN back, + clear its state, and stay usable for the next transaction. + """ + hass, _registry = make_hass() + database = NestQuestDatabase(hass) + await database.open(tmp_path / "cancel-begin.db") + await database.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)") + await database.execute("INSERT INTO t (name) VALUES (?)", ("before",)) + + begin_started = asyncio.Event() + release_begin = asyncio.Event() + real_add = hass.async_add_executor_job + + async def _gated_add_executor_job(fn, *args, **kwargs): + if getattr(fn, "__name__", "") == "_begin": + begin_started.set() + await release_begin.wait() + return await real_add(fn, *args, **kwargs) + + hass.async_add_executor_job = _gated_add_executor_job + + async def _txn() -> None: + async with database.transaction(): + await database.execute("INSERT INTO t (name) VALUES (?)", ("kept",)) + + txn_task = asyncio.create_task(_txn()) + await begin_started.wait() + txn_task.cancel() + await asyncio.sleep(0.05) + assert database._in_transaction is True, ( + "ownership was released while the cancelled BEGIN was still in flight" + ) + release_begin.set() + with pytest.raises(asyncio.CancelledError): + await txn_task + + assert database._in_transaction is False + assert database._tx_task is None + assert database._conn is not None + assert database._conn.in_transaction is False + + hass.async_add_executor_job = real_add + assert await database.fetch_all("SELECT name FROM t ORDER BY id") == [ + ("before",) + ] + async with database.transaction(): + await database.execute("INSERT INTO t (name) VALUES (?)", ("kept",)) + assert await database.fetch_all("SELECT name FROM t ORDER BY id") == [ + ("before",), + ("kept",), + ] + await database.close() + + +async def test_cancel_during_commit_settles_then_propagates_cancellation( + tmp_path, +) -> None: + """NQ-PR12-007: cancel during COMMIT must not mask the cancellation. + + COMMIT is gated; the task is cancelled while it is in flight and the + worker then completes it anyway. The cancellation must propagate as + CancelledError (not be converted by a doomed ROLLBACK of an already + committed transaction), the connection must end outside any + transaction, and the wrapper must stay usable. + """ + hass, _registry = make_hass() + database = NestQuestDatabase(hass) + await database.open(tmp_path / "cancel-commit.db") + await database.execute("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)") + + commit_started = asyncio.Event() + release_commit = asyncio.Event() + real_add = hass.async_add_executor_job + + async def _gated_add_executor_job(fn, *args, **kwargs): + if getattr(fn, "__name__", "") == "_commit": + commit_started.set() + await release_commit.wait() + return await real_add(fn, *args, **kwargs) + + hass.async_add_executor_job = _gated_add_executor_job + + async def _txn() -> None: + async with database.transaction(): + await database.execute("INSERT INTO t (name) VALUES (?)", ("kept",)) + + txn_task = asyncio.create_task(_txn()) + await commit_started.wait() + txn_task.cancel() + await asyncio.sleep(0.05) + assert database._in_transaction is True, ( + "ownership was released while the cancelled COMMIT was still in flight" + ) + release_commit.set() + with pytest.raises(asyncio.CancelledError): + await txn_task + + assert database._in_transaction is False + assert database._tx_task is None + assert database._conn is not None + assert database._conn.in_transaction is False + + hass.async_add_executor_job = real_add + assert await database.fetch_all("SELECT name FROM t") == [("kept",)] + async with database.transaction(): + await database.execute("INSERT INTO t (name) VALUES (?)", ("again",)) + assert await database.fetch_all("SELECT name FROM t ORDER BY id") == [ + ("kept",), + ("again",), + ] + await database.close() + + +async def test_commit_and_rollback_failures_are_both_preserved(tmp_path) -> None: + """NQ-PR12-008: both errors survive when COMMIT and ROLLBACK both fail. + + The COMMIT executor call is forced to fail, and the follow-up + cleanup ROLLBACK is forced to fail as well. The raised error must + keep both original exceptions inspectable instead of discarding the + commit failure behind the rollback failure. + """ + hass, _registry = make_hass() + database = NestQuestDatabase(hass) + await database.open(tmp_path / "both-fail.db") + await database.execute("CREATE TABLE t (id INTEGER PRIMARY KEY)") + + commit_failure = sqlite3.OperationalError("commit exploded") + rollback_failure = sqlite3.OperationalError("rollback exploded") + real_add = hass.async_add_executor_job + + async def _failing_add_executor_job(fn, *args, **kwargs): + name = getattr(fn, "__name__", "") + if name == "_commit": + raise commit_failure + if name == "_cleanup": + raise rollback_failure + return await real_add(fn, *args, **kwargs) + + hass.async_add_executor_job = _failing_add_executor_job + + with pytest.raises(BaseExceptionGroup) as exc_info: + async with database.transaction(): + pass + + assert commit_failure in exc_info.value.exceptions + assert rollback_failure in exc_info.value.exceptions + assert database._in_transaction is False + assert database._tx_task is None + await database.close() + + async def test_transaction_persists_across_connections(tmp_path) -> None: """A committed transaction is visible to a brand-new connection.""" hass, _registry = make_hass()