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._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.