mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-28 01:31:28 +00:00
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:
parent
6bda95d3bd
commit
a03bde4701
2 changed files with 82 additions and 32 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue