feat: add async sqlite connection wrapper with executor-backed queries

2026-09-13 20:20:31 -04:00
parent 221222275a
commit ae7a064d63
5 changed files with 570 additions and 2 deletions
@@ -0,0 +1,143 @@
+"""Async SQLite connection wrapper for the NestQuest integration."""
+from __future__ import annotations
+
+import sqlite3
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+
+from homeassistant.core import HomeAssistant
+
+
+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.
+    """
+
+    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
+
+    @property
+    def connected(self) -> bool:
+        """Return True while the underlying connection is open."""
+        return self._conn is not None
+
+    @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
+
+    @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
+
+    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
+
+        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
+
+        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()
+
+        def _execute() -> sqlite3.Cursor:
+            return conn.execute(sql, parameters)
+
+        self._cursor = await self._hass.async_add_executor_job(_execute)
+        return self._cursor
+
+    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()
+
+        def _execute_many() -> sqlite3.Cursor:
+            return conn.executemany(sql, parameters)
+
+        self._cursor = await self._hass.async_add_executor_job(_execute_many)
+        return self._cursor
+
+    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)
+
+        def _fetch_one() -> tuple | None:
+            return cursor.fetchone()
+
+        return await self._hass.async_add_executor_job(_fetch_one)
+
+    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)
+
+        def _fetch_all() -> list[tuple]:
+            return cursor.fetchall()
+
+        return await self._hass.async_add_executor_job(_fetch_all)
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncIterator[None]:
+        """Run a block inside BEGIN/COMMIT, rolling back on any exception."""
+        conn = self._require_conn()
+
+        def _begin() -> None:
+            conn.execute("BEGIN")
+
+        await self._hass.async_add_executor_job(_begin)
+        try:
+            yield
+        except BaseException:
+
+            def _rollback() -> None:
+                conn.execute("ROLLBACK")
+
+            await self._hass.async_add_executor_job(_rollback)
+            raise
+
+        def _commit() -> None:
+            conn.execute("COMMIT")
+
+        await self._hass.async_add_executor_job(_commit)
+
+    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
+
+        def _close() -> None:
+            conn.close()
+
+        await self._hass.async_add_executor_job(_close)
+
+    def _require_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            raise RuntimeError(
+                "The NestQuest database connection is closed; call open() first"
+            )
+        return self._conn