fix: close abandoned connection when open is cancelled after settle

This commit is contained in:
Joshua King
2026-09-13 21:33:43 -04:00
parent 97c8040e05
commit c646446f33
2 changed files with 102 additions and 5 deletions
+37 -5
View File
@@ -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.
+65
View File
@@ -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.