fix(db): tighten aiosqlite terminate-shim install scope and worker wake-up

Address PR review findings on the original shim:

1. Source-shape gating (avoid shadowing future upstream fixes)
   The install now inspects the original implementation source and
   only patches when it actually references `self._connection.stop` —
   the buggy pattern this shim exists to fix. If a newer SQLAlchemy
   ships a working `_terminate_force_close` we leave it alone.

2. Wake the aiosqlite worker, don't just flag it
   In aiosqlite ≥0.20 the worker thread blocks on `_tx.get()`. Setting
   `_running = False` alone never unblocks it — the queue needs a
   sentinel push for the loop to observe the flag. The replacement
   now puts `(None, None)` onto `_tx` first and then flips
   `_running`, so the worker exits promptly instead of relying on a
   subsequent unrelated queue item or GC.

3. Conditional install (no global mutation on non-sqlite deployments)
   Move `_install_aiosqlite_compat()` inside the
   `if 'sqlite' in ASYNC_SQLALCHEMY_DATABASE_URL` branch in
   `internal/db.py`, so PostgreSQL / MySQL / etc. deployments get no
   SQLAlchemy monkey-patch at all.

https://claude.ai/code/session_01JSr4NZSskEUQvoJnavVXh8
This commit is contained in:
Claude 2026-04-14 12:41:35 +00:00
parent 6bda95d3bd
commit a03bde4701
No known key found for this signature in database
2 changed files with 82 additions and 32 deletions

View file

@ -30,19 +30,34 @@ worker on the next iteration.
Fix
---
Replace `_terminate_force_close` with one that understands both the
old `Connection.stop()` API and the modern internal `_running` flag.
Setting `_running = False` causes the aiosqlite worker thread to
break out of its loop on the next tick, which is the same end-state
the original code was after.
old `Connection.stop()` API and the modern post-0.20 internals (a
private worker thread fed by an `_tx` queue, gated by a `_running`
flag). The replacement is best-effort: the connection record has
already been removed from the pool by the time this is called, and
aiosqlite's own GC will release the underlying sqlite3 handle. We
nudge the worker (sentinel + `_running = False`) so it doesn't linger
on the queue indefinitely, but we do not raise on failure.
Apply this patch once, at import time, before any async engine is
created. Idempotent and safe to import on systems where the affected
SQLAlchemy/aiosqlite versions are not in use — it bails out silently
if the symbol is missing.
created. To keep the blast radius small the install is:
* a no-op if the SQLAlchemy build doesn't ship the affected
`_terminate_force_close` symbol (older or upstream-fixed
versions);
* a no-op if the upstream implementation has already moved off the
`self._connection.stop` reference — detected by inspecting the
original source — so a future SQLAlchemy fix isn't shadowed by
this shim;
* idempotent so repeated imports don't stack patches.
The caller in `internal.db` only invokes `install()` when the runtime
async URL is sqlite-based, so PostgreSQL / MySQL / etc. deployments
get no global mutation at all.
"""
from __future__ import annotations
import inspect
import logging
log = logging.getLogger(__name__)
@ -61,15 +76,35 @@ def install() -> None:
# upstream rewrite that no longer needs this shim.
return
if getattr(target_cls._terminate_force_close, '__open_webui_patched__', False):
original = target_cls._terminate_force_close
if getattr(original, '__open_webui_patched__', False):
return # Idempotent — already applied.
# Only patch the specific buggy implementation. If upstream has
# changed how `_terminate_force_close` is implemented (no longer
# touching `self._connection.stop`) defer to whatever they ship.
try:
original_source = inspect.getsource(original)
except (OSError, TypeError):
original_source = ''
if 'self._connection.stop' not in original_source:
return
def _terminate_force_close(self) -> None:
"""Best-effort force close of an aiosqlite connection.
The pool has already removed the connection record by the time
this is called; aiosqlite's own GC will release the underlying
sqlite3 handle either way. We try to nudge the worker so it
exits promptly on both the pre-0.18 and post-0.20 internals,
but we never raise: there is nothing useful for the caller to
do with a failure here.
"""
conn = getattr(self, '_connection', None)
if conn is None:
return
# aiosqlite ≤ 0.18 — original API.
# aiosqlite ≤ 0.18 — original public API.
stop = getattr(conn, 'stop', None)
if callable(stop):
try:
@ -82,29 +117,40 @@ def install() -> None:
)
return
# aiosqlite ≥ 0.20 — the worker thread observes its own
# `_running` flag and exits on the next loop tick, releasing
# the underlying sqlite3 connection.
# aiosqlite ≥ 0.20 — the worker thread blocks on `_tx.get()`
# waiting for callables. Setting `_running = False` alone does
# not wake it; the queue needs a sentinel push so the loop
# observes the flag on the next iteration.
nudged = False
tx = getattr(conn, '_tx', None)
if tx is not None:
try:
tx.put_nowait((None, None))
nudged = True
except Exception:
log.debug(
'aiosqlite worker queue rejected the termination '
'sentinel; relying on garbage collection.',
exc_info=True,
)
if hasattr(conn, '_running'):
try:
conn._running = False
nudged = True
except Exception:
log.debug(
'Could not flip aiosqlite Connection._running during '
'force-close; connection will be cleaned up by garbage '
'collection.',
'force-close.',
exc_info=True,
)
return
# Unknown aiosqlite internals — fall back to leaving the
# connection to garbage collection. Do *not* re-raise: the pool
# has already removed the connection record and there is
# nothing useful for the caller to do.
log.debug(
'aiosqlite Connection has neither stop() nor _running; relying '
'on garbage collection to release the underlying sqlite3 handle.'
)
if not nudged:
log.debug(
'aiosqlite Connection has neither stop(), _tx nor _running; '
'relying on garbage collection to release the underlying '
'sqlite3 handle.'
)
_terminate_force_close.__open_webui_patched__ = True # type: ignore[attr-defined]
target_cls._terminate_force_close = _terminate_force_close

View file

@ -205,18 +205,22 @@ get_db = contextmanager(get_session)
# ASYNC ENGINE (used for ALL runtime database operations)
# ============================================================
# Patch SQLAlchemy's aiosqlite connector before any async engine is
# created. Without this, every cancelled aiosqlite query produces a
# multi-page `terminate_force_close() not implemented` ERROR traceback
# because SQLAlchemy 2.0.x's shim references `Connection.stop`, which
# aiosqlite removed in 0.20+. See `_aiosqlite_compat` for details.
from open_webui.internal._aiosqlite_compat import install as _install_aiosqlite_compat
_install_aiosqlite_compat()
ASYNC_SQLALCHEMY_DATABASE_URL = _make_async_url(SQLALCHEMY_DATABASE_URL)
if 'sqlite' in ASYNC_SQLALCHEMY_DATABASE_URL:
# Patch SQLAlchemy's aiosqlite connector before the engine — and
# therefore any pooled connection — is created. Without this, every
# cancelled aiosqlite query produces a multi-page
# `terminate_force_close() not implemented` ERROR traceback because
# SQLAlchemy 2.0.x's shim references `Connection.stop`, which
# aiosqlite removed in 0.20+. The install is a no-op when the
# affected upstream code is absent or already fixed; it is gated
# behind the sqlite URL check so non-sqlite deployments get no
# global SQLAlchemy mutation. See `_aiosqlite_compat` for details.
from open_webui.internal._aiosqlite_compat import install as _install_aiosqlite_compat
_install_aiosqlite_compat()
async_engine = create_async_engine(
ASYNC_SQLALCHEMY_DATABASE_URL,
connect_args={'check_same_thread': False},