overseer
fb3c6601c1
HACS and the manual install copy only custom_components/nestquest, so the integration shipped with a top-level 'core' package that was never deployed and failed with ModuleNotFoundError: No module named 'core'. Move core/ into custom_components/nestquest/core/, repoint the re-export shims to relative '.core.X' imports, and update tests (import paths, monkeypatch targets, caplog logger) to the new 'custom_components.nestquest.core.X' location. The integration is now self-contained: nothing under custom_components/nestquest reaches outside its own tree, the migrations logger becomes custom_components.nestquest.core.migrations, and the SQL-leak guard allowlists point at the new core/ paths. uv run pytest -q: 1496 passed.
436 lines
17 KiB
Python
436 lines
17 KiB
Python
"""Async SQLite connection wrapper for the NestQuest integration."""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import sqlite3
|
|
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
#: Executor contract: schedule a synchronous callable on a worker
|
|
#: thread/futures pool and await its result. Mirrors the shape of
|
|
#: ``hass.async_add_executor_job`` so the integration wires it in with a
|
|
#: bare method reference, but this package never imports Home Assistant.
|
|
Executor = Callable[..., Awaitable[T]]
|
|
|
|
|
|
@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 a worker through the ``executor``
|
|
callable passed at construction (mirroring the contract of
|
|
``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.
|
|
|
|
Every connection-affecting executor job (connect, pragmas, statements,
|
|
transaction lifecycle and close) is awaited through the settlement
|
|
helper :meth:`_await_settled`: if the awaiting task is cancelled while
|
|
a job is in flight, the job still settles on the worker thread before
|
|
the cancellation propagates, so the wrapper lock is never released
|
|
while the executor thread is mid-job.
|
|
"""
|
|
|
|
def __init__(self, executor: Callable[..., Awaitable[T]]) -> None:
|
|
"""Initialize the wrapper around an executor callable.
|
|
|
|
``executor`` schedules a synchronous callable on a worker and
|
|
returns an awaitable of its result — the contract of
|
|
``hass.async_add_executor_job``. The integration passes that
|
|
method directly; tests pass a helper backed by
|
|
:func:`asyncio.to_thread` or an inline runner.
|
|
"""
|
|
self._executor: Callable[..., Awaitable[Any]] = executor
|
|
self._conn: sqlite3.Connection | 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:
|
|
"""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._rowcount
|
|
|
|
@property
|
|
def lastrowid(self) -> int | None:
|
|
"""Row id of the most recent insert, or None when unset."""
|
|
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.
|
|
|
|
Cancellation-safe: if this call is cancelled while the connect job
|
|
is in flight, the job still settles (see :meth:`_await_settled`).
|
|
A connection the settled job created is never adopted by the
|
|
cancelled call; it is closed through a settled executor job instead
|
|
of being abandoned to garbage collection, and the cancellation
|
|
propagates afterwards, so a follow-up ``open()`` starts from a
|
|
clean state.
|
|
"""
|
|
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
|
|
)
|
|
try:
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
except BaseException:
|
|
conn.close()
|
|
raise
|
|
return conn
|
|
|
|
settled, cancelled = await self._await_settled_state(
|
|
self._executor(_open)
|
|
)
|
|
if cancelled:
|
|
if isinstance(settled, sqlite3.Connection):
|
|
|
|
def _close_abandoned() -> None:
|
|
settled.close()
|
|
|
|
await self._await_settled(
|
|
self._executor(_close_abandoned)
|
|
)
|
|
raise asyncio.CancelledError()
|
|
self._conn = settled
|
|
return self
|
|
|
|
async def execute(
|
|
self, sql: str, parameters: Sequence[Any] = ()
|
|
) -> ExecutionResult:
|
|
"""Run one statement in the executor and return its plain result."""
|
|
async with self._serialized():
|
|
conn = self._require_conn()
|
|
|
|
def _execute() -> ExecutionResult:
|
|
cursor = conn.execute(sql, parameters)
|
|
return ExecutionResult(cursor.rowcount, cursor.lastrowid)
|
|
|
|
result = await self._await_settled(
|
|
self._executor(_execute)
|
|
)
|
|
self._rowcount = result.rowcount
|
|
self._lastrowid = result.lastrowid
|
|
return result
|
|
|
|
async def execute_many(
|
|
self, sql: str, parameters: Sequence[Sequence[Any]]
|
|
) -> ExecutionResult:
|
|
"""Run executemany in the executor and return its plain result."""
|
|
async with self._serialized():
|
|
conn = self._require_conn()
|
|
|
|
def _execute_many() -> ExecutionResult:
|
|
cursor = conn.executemany(sql, parameters)
|
|
return ExecutionResult(cursor.rowcount, cursor.lastrowid)
|
|
|
|
result = await self._await_settled(
|
|
self._executor(_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."""
|
|
async with self._serialized():
|
|
conn = self._require_conn()
|
|
|
|
def _fetch_one() -> tuple[tuple | None, int, int | None]:
|
|
cursor = conn.execute(sql, parameters)
|
|
return cursor.fetchone(), cursor.rowcount, cursor.lastrowid
|
|
|
|
row, rowcount, lastrowid = await self._await_settled(
|
|
self._executor(_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."""
|
|
async with self._serialized():
|
|
conn = self._require_conn()
|
|
|
|
def _fetch_all() -> tuple[list[tuple], int, int | None]:
|
|
cursor = conn.execute(sql, parameters)
|
|
return cursor.fetchall(), cursor.rowcount, cursor.lastrowid
|
|
|
|
rows, rowcount, lastrowid = await self._await_settled(
|
|
self._executor(_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.
|
|
|
|
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. A
|
|
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(
|
|
"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:
|
|
|
|
def _begin() -> None:
|
|
conn.execute("BEGIN")
|
|
|
|
await self._settle_lifecycle(
|
|
conn,
|
|
self._executor(_begin),
|
|
"BEGIN did not settle cleanly and the follow-up "
|
|
"cleanup also raised; the transaction may need manual "
|
|
"cleanup",
|
|
)
|
|
try:
|
|
yield
|
|
except BaseException:
|
|
|
|
def _rollback() -> None:
|
|
conn.execute("ROLLBACK")
|
|
|
|
await self._settle_lifecycle(
|
|
conn,
|
|
self._executor(_rollback),
|
|
"ROLLBACK failed and the follow-up cleanup also "
|
|
"raised; the transaction may need manual cleanup",
|
|
)
|
|
raise
|
|
|
|
def _commit() -> None:
|
|
conn.execute("COMMIT")
|
|
|
|
await self._settle_lifecycle(
|
|
conn,
|
|
self._executor(_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._executor(_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.
|
|
"""
|
|
result, cancelled = await self._await_settled_state(job)
|
|
if cancelled:
|
|
raise asyncio.CancelledError()
|
|
return result
|
|
|
|
async def _await_settled_state(self, job: Awaitable[Any]) -> tuple[Any, bool]:
|
|
"""Await ``job`` to settlement and report ``(result, cancelled)``.
|
|
|
|
Same settlement guarantees as :meth:`_await_settled`, but returns
|
|
the settled outcome together with whether the awaiting task was
|
|
cancelled, so callers that must act on a settled result before
|
|
re-raising cancellation (e.g. closing an abandoned connection)
|
|
can do so without losing the result.
|
|
"""
|
|
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
|
|
return result, cancelled
|
|
|
|
async def close(self) -> None:
|
|
"""Close the connection in the executor; safe to call repeatedly.
|
|
|
|
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
|
|
|
|
def _close() -> None:
|
|
conn.close()
|
|
|
|
await self._await_settled(self._executor(_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 |