fix: make db wrapper executor-safe, serialized, and leak-proof

2026-09-13 21:07:42 -04:00
parent ae7a064d63
commit a56a7e3d37
4 changed files with 525 additions and 67 deletions
@@ -50,11 +50,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
        await listener_hass.config_entries.async_reload(listener_entry.entry_id)

    database = await NestQuestDatabase(hass).open(await async_get_db_path(hass))
+    try:
+        remove_update_listener = entry.add_update_listener(_async_update_listener)
+    except BaseException:
+        await database.close()
+        raise
    runtime_data = NestQuestRuntimeData(
        entry_id=entry.entry_id,
        options=dict(entry.options),
        database=database,
-        remove_update_listener=entry.add_update_listener(_async_update_listener),
+        remove_update_listener=remove_update_listener,
    )
    entry.runtime_data = runtime_data
    hass.data[DOMAIN][entry.entry_id] = runtime_data
@@ -71,13 +76,22 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
        runtime_data = getattr(entry, "runtime_data", None)
    if runtime_data is not None:
        remove_update_listener = getattr(runtime_data, "remove_update_listener", None)
-        if remove_update_listener is not None:
-            result = remove_update_listener()
-            if inspect.isawaitable(result):
-                await result
-        database = getattr(runtime_data, "database", None)
-        if database is not None:
-            await database.close()
+        unload_error: BaseException | None = None
+        try:
+            if remove_update_listener is not None:
+                result = remove_update_listener()
+                if inspect.isawaitable(result):
+                    await result
+        except BaseException as err:
+            unload_error = err
+        finally:
+            database = getattr(runtime_data, "database", None)
+            if database is not None:
+                await database.close()
+        if unload_error is not None:
+            if getattr(entry, "runtime_data", None) is not None:
+                entry.runtime_data = None
+            raise unload_error
    if getattr(entry, "runtime_data", None) is not None:
        entry.runtime_data = None
    return True
@@ -1,27 +1,72 @@
 """Async SQLite connection wrapper for the NestQuest integration."""
 from __future__ import annotations

+import asyncio
 import sqlite3
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any

 from homeassistant.core import HomeAssistant


+@dataclass(frozen=True)
+class ExecutionResult:
+    """Plain-data outcome of one statement.
+
+    ``rowcount`` and ``lastrowid`` mirror the corresponding
+    :class:`sqlite3.Cursor` attributes; both are captured on the executor
+    thread so no sqlite3 object ever crosses back to the event loop.
+    """
+
+    rowcount: int
+    lastrowid: int | None
+
+
 class NestQuestDatabase:
    """Async wrapper around a single SQLite connection.

    Every sqlite3 call is delegated to the Home Assistant executor through
-    ``hass.async_add_executor_job`` so the event loop is never blocked.
+    ``hass.async_add_executor_job`` so the event loop is never blocked, and
+    only plain Python data (row tuples and :class:`ExecutionResult`) is
+    returned to the event loop.
+
+    Every operation is serialized through one asyncio.Lock per wrapper.  A
+    ``transaction()`` holds that lock for its whole BEGIN..COMMIT span, so
+    interleaved coroutines cannot corrupt the transaction and concurrent
+    BEGINs are impossible; statements issued while a transaction is open
+    queue until it completes.  Concurrent transactions on one wrapper are
+    not supported: a second ``transaction()`` entered while one is active
+    raises :class:`RuntimeError`.  ``close()`` serializes on the same lock
+    and is idempotent.
    """

    def __init__(self, hass: HomeAssistant) -> None:
        """Initialize the wrapper around a Home Assistant instance."""
        self._hass = hass
        self._conn: sqlite3.Connection | None = None
-        self._cursor: sqlite3.Cursor | None = None
+        self._rowcount: int = -1
+        self._lastrowid: int | None = None
+        self._lock: asyncio.Lock | None = None
+        self._lock_loop: asyncio.AbstractEventLoop | None = None
+        self._in_transaction = False
+        self._tx_task: asyncio.Task | None = None
+
+    def _get_lock(self) -> asyncio.Lock:
+        """Return the per-wrapper lock, (re)bound to the running loop.
+
+        HA drives every wrapper call from one event loop, so the lock is
+        created once and reused; re-binding only when the running loop
+        differs keeps the wrapper usable from plain synchronous drivers
+        that run each call on a fresh loop.
+        """
+        loop = asyncio.get_running_loop()
+        if self._lock is None or self._lock_loop is not loop:
+            self._lock = asyncio.Lock()
+            self._lock_loop = loop
+        return self._lock

    @property
    def connected(self) -> bool:
@@ -31,109 +76,189 @@ class NestQuestDatabase:
    @property
    def rowcount(self) -> int:
        """Row count of the most recent statement, or -1 when unset."""
-        return self._cursor.rowcount if self._cursor is not None else -1
+        return self._rowcount

    @property
    def lastrowid(self) -> int | None:
        """Row id of the most recent insert, or None when unset."""
-        return self._cursor.lastrowid if self._cursor is not None else None
+        return self._lastrowid
+
+    @asynccontextmanager
+    async def _serialized(self) -> AsyncIterator[None]:
+        """Hold the wrapper lock unless the caller owns the active transaction.
+
+        A transaction holds the lock for its whole BEGIN..COMMIT span;
+        statements issued by the transaction's own task are already
+        serialized by it and must not re-acquire the non-reentrant lock.
+        Any other task queues on the lock until the transaction ends.
+        """
+        if (
+            self._in_transaction
+            and self._tx_task is not None
+            and asyncio.current_task() is self._tx_task
+        ):
+            yield
+            return
+        async with self._get_lock():
+            yield

    async def open(self, path: Path | str) -> NestQuestDatabase:
        """Open the SQLite file with WAL journal mode and foreign keys on."""
        if self._conn is not None:
            return self
+        async with self._get_lock():
+            if self._conn is not None:
+                return self

-        def _open() -> sqlite3.Connection:
-            conn = sqlite3.connect(
-                path, check_same_thread=False, isolation_level=None
-            )
-            conn.execute("PRAGMA journal_mode=WAL")
-            conn.execute("PRAGMA foreign_keys=ON")
-            return conn
+            def _open() -> sqlite3.Connection:
+                conn = sqlite3.connect(
+                    path, check_same_thread=False, isolation_level=None
+                )
+                try:
+                    conn.execute("PRAGMA journal_mode=WAL")
+                    conn.execute("PRAGMA foreign_keys=ON")
+                except BaseException:
+                    conn.close()
+                    raise
+                return conn

-        self._conn = await self._hass.async_add_executor_job(_open)
+            self._conn = await self._hass.async_add_executor_job(_open)
        return self

    async def execute(
        self, sql: str, parameters: Sequence[Any] = ()
-    ) -> sqlite3.Cursor:
-        """Run one statement in the executor and return its cursor."""
-        conn = self._require_conn()
+    ) -> ExecutionResult:
+        """Run one statement in the executor and return its plain result."""
+        async with self._serialized():
+            conn = self._require_conn()

-        def _execute() -> sqlite3.Cursor:
-            return conn.execute(sql, parameters)
+            def _execute() -> ExecutionResult:
+                cursor = conn.execute(sql, parameters)
+                return ExecutionResult(cursor.rowcount, cursor.lastrowid)

-        self._cursor = await self._hass.async_add_executor_job(_execute)
-        return self._cursor
+            result = await self._hass.async_add_executor_job(_execute)
+            self._rowcount = result.rowcount
+            self._lastrowid = result.lastrowid
+            return result

    async def execute_many(
        self, sql: str, parameters: Sequence[Sequence[Any]]
-    ) -> sqlite3.Cursor:
-        """Run executemany in the executor and return its cursor."""
-        conn = self._require_conn()
+    ) -> ExecutionResult:
+        """Run executemany in the executor and return its plain result."""
+        async with self._serialized():
+            conn = self._require_conn()

-        def _execute_many() -> sqlite3.Cursor:
-            return conn.executemany(sql, parameters)
+            def _execute_many() -> ExecutionResult:
+                cursor = conn.executemany(sql, parameters)
+                return ExecutionResult(cursor.rowcount, cursor.lastrowid)

-        self._cursor = await self._hass.async_add_executor_job(_execute_many)
-        return self._cursor
+            result = await self._hass.async_add_executor_job(_execute_many)
+            self._rowcount = result.rowcount
+            self._lastrowid = result.lastrowid
+            return result

    async def fetch_one(
        self, sql: str, parameters: Sequence[Any] = ()
    ) -> tuple | None:
        """Run a query in the executor and return its first row, or None."""
-        cursor = await self.execute(sql, parameters)
+        async with self._serialized():
+            conn = self._require_conn()

-        def _fetch_one() -> tuple | None:
-            return cursor.fetchone()
+            def _fetch_one() -> tuple[tuple | None, int, int | None]:
+                cursor = conn.execute(sql, parameters)
+                return cursor.fetchone(), cursor.rowcount, cursor.lastrowid

-        return await self._hass.async_add_executor_job(_fetch_one)
+            row, rowcount, lastrowid = await self._hass.async_add_executor_job(
+                _fetch_one
+            )
+            self._rowcount = rowcount
+            self._lastrowid = lastrowid
+            return row

    async def fetch_all(
        self, sql: str, parameters: Sequence[Any] = ()
    ) -> list[tuple]:
        """Run a query in the executor and return all of its rows."""
-        cursor = await self.execute(sql, parameters)
+        async with self._serialized():
+            conn = self._require_conn()

-        def _fetch_all() -> list[tuple]:
-            return cursor.fetchall()
+            def _fetch_all() -> tuple[list[tuple], int, int | None]:
+                cursor = conn.execute(sql, parameters)
+                return cursor.fetchall(), cursor.rowcount, cursor.lastrowid

-        return await self._hass.async_add_executor_job(_fetch_all)
+            rows, rowcount, lastrowid = await self._hass.async_add_executor_job(
+                _fetch_all
+            )
+            self._rowcount = rowcount
+            self._lastrowid = lastrowid
+            return rows

    @asynccontextmanager
    async def transaction(self) -> AsyncIterator[None]:
-        """Run a block inside BEGIN/COMMIT, rolling back on any exception."""
-        conn = self._require_conn()
+        """Run a block inside BEGIN/COMMIT, rolling back on any exception.

-        def _begin() -> None:
-            conn.execute("BEGIN")
+        The wrapper lock is held for the entire BEGIN..COMMIT span, so no
+        other statement can interleave inside the transaction.  Entering a
+        second transaction while one is active raises RuntimeError.
+        """
+        if self._in_transaction:
+            raise RuntimeError(
+                "A NestQuest transaction is already active on this "
+                "connection; concurrent transactions are not supported"
+            )
+        async with self._get_lock():
+            conn = self._require_conn()
+            self._in_transaction = True
+            self._tx_task = asyncio.current_task()
+            try:

-        await self._hass.async_add_executor_job(_begin)
-        try:
-            yield
-        except BaseException:
+                def _begin() -> None:
+                    conn.execute("BEGIN")

-            def _rollback() -> None:
-                conn.execute("ROLLBACK")
+                await self._hass.async_add_executor_job(_begin)
+                try:
+                    yield
+                except BaseException:

-            await self._hass.async_add_executor_job(_rollback)
-            raise
+                    def _rollback() -> None:
+                        conn.execute("ROLLBACK")

-        def _commit() -> None:
-            conn.execute("COMMIT")
+                    await self._hass.async_add_executor_job(_rollback)
+                    raise

-        await self._hass.async_add_executor_job(_commit)
+                def _commit() -> None:
+                    conn.execute("COMMIT")
+
+                await self._hass.async_add_executor_job(_commit)
+            finally:
+                self._in_transaction = False
+                self._tx_task = None

    async def close(self) -> None:
-        """Close the connection in the executor; safe to call repeatedly."""
-        conn, self._conn, self._cursor = self._conn, None, None
-        if conn is None:
-            return
+        """Close the connection in the executor; safe to call repeatedly.

-        def _close() -> None:
-            conn.close()
+        Serialized on the wrapper lock, so an in-flight operation or open
+        transaction always completes before the connection is closed.
+        """
+        if (
+            self._in_transaction
+            and self._tx_task is not None
+            and asyncio.current_task() is self._tx_task
+        ):
+            raise RuntimeError(
+                "Cannot close the NestQuest database from inside its own "
+                "transaction; exit the transaction first"
+            )
+        async with self._get_lock():
+            conn, self._conn = self._conn, None
+            self._rowcount, self._lastrowid = -1, None
+            if conn is None:
+                return

-        await self._hass.async_add_executor_job(_close)
+            def _close() -> None:
+                conn.close()
+
+            await self._hass.async_add_executor_job(_close)

    def _require_conn(self) -> sqlite3.Connection:
        if self._conn is None: