fix: make transaction lifecycle cancellation-safe and preserve commit+rollback errors

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._settle_lifecycle(
-                    await self._hass.async_add_executor_job(_commit)
+                    conn,
-                except BaseException as commit_error:
+                    self._hass.async_add_executor_job(_commit),
-
+                    "COMMIT failed and the follow-up cleanup also raised; "
-                    def _rollback() -> None:
+                    "the transaction may need manual cleanup",
-                        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
             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.
 

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