diff --git a/custom_components/nestquest/db.py b/custom_components/nestquest/db.py index bd785a4..c05af6b 100644 --- a/custom_components/nestquest/db.py +++ b/custom_components/nestquest/db.py @@ -110,7 +110,16 @@ class NestQuestDatabase: yield async def open(self, path: Path | str) -> NestQuestDatabase: - """Open the SQLite file with WAL journal mode and foreign keys on.""" + """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(): @@ -129,9 +138,20 @@ class NestQuestDatabase: raise return conn - self._conn = await self._await_settled( + settled, cancelled = await self._await_settled_state( self._hass.async_add_executor_job(_open) ) + if cancelled: + if isinstance(settled, sqlite3.Connection): + + def _close_abandoned() -> None: + settled.close() + + await self._await_settled( + self._hass.async_add_executor_job(_close_abandoned) + ) + raise asyncio.CancelledError() + self._conn = settled return self async def execute( @@ -336,6 +356,20 @@ class NestQuestDatabase: 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(): @@ -352,9 +386,7 @@ class NestQuestDatabase: [asyncio.CancelledError(), job_error], ) from job_error raise - if cancelled: - raise asyncio.CancelledError() - return result + return result, cancelled async def close(self) -> None: """Close the connection in the executor; safe to call repeatedly. diff --git a/tests/test_db.py b/tests/test_db.py index 8aceb8c..7640c57 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -406,6 +406,71 @@ async def test_commit_failure_rolls_back_and_next_transaction_works( await database.close() +async def test_cancel_during_open_closes_abandoned_connection(tmp_path) -> None: + """NQ-PR12-010: cancel during open() must close the settled connection. + + The _open executor job is gated: the open task is cancelled while the + connect is still in flight, then the worker is allowed to complete it + anyway. The settled connection must never be adopted, must be closed + through a settled executor-backed close before the cancellation + propagates, and a follow-up open() must succeed cleanly. + """ + hass, _registry = make_hass() + database = NestQuestDatabase(hass) + + open_started = asyncio.Event() + release_open = asyncio.Event() + jobs: list[str] = [] + captured: dict[str, Any] = {} + real_add = hass.async_add_executor_job + + async def _gated_add_executor_job(fn, *args, **kwargs): + name = getattr(fn, "__name__", "") + jobs.append(name) + if name == "_open": + open_started.set() + await release_open.wait() + result = await real_add(fn, *args, **kwargs) + if name == "_open": + captured["conn"] = result + return result + + hass.async_add_executor_job = _gated_add_executor_job + + open_task = asyncio.create_task(database.open(tmp_path / "cancel-open.db")) + await open_started.wait() + open_task.cancel() + await asyncio.sleep(0.05) + assert database._conn is None, ( + "the cancelled open adopted a connection that was still settling" + ) + release_open.set() + with pytest.raises(asyncio.CancelledError): + await open_task + + assert database.connected is False + assert "_close_abandoned" in jobs, ( + "the abandoned connection was not closed via an executor job" + ) + conn = captured.get("conn") + assert conn is not None + + def _probe() -> str: + try: + conn.execute("SELECT 1") + except sqlite3.ProgrammingError: + return "closed" + return "open" + + assert await real_add(_probe) == "closed" + + hass.async_add_executor_job = real_add + await database.open(tmp_path / "cancel-open.db") + assert database.connected is True + assert await database.fetch_one("SELECT 1") == (1,) + 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.