fix: cache app-password storage to avoid per-request Alembic upgrade race

get_user_client_basic_auth built a fresh RefreshTokenStorage and ran
storage.initialize() — a full Alembic `upgrade` in a worker thread — on every
call. Once the /api/v1 search endpoints (unified_search, vector_search) and the
chunk-context endpoint were wired to use it, concurrent requests ran concurrent
Alembic upgrades, which race on Alembic's non-thread-safe module-global
EnvironmentContext proxy and intermittently raise `KeyError: 'script'` →
HTTP 500 (seen on multi-user-basic/nc31; nc32 got lucky).

Cache one process-wide, already-initialized storage instance behind a lazily
created anyio.Lock so the one-time migration runs exactly once and never races.
Callers passing an explicit `storage` are unaffected. Also removes a redundant
per-request migration from the search hot path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chris Coutinho
2026-05-29 17:37:36 +02:00
co-authored by Claude Opus 4.8
parent 350358b802
commit 063bd06ba6
+36 -3
View File
@@ -45,6 +45,38 @@ class NotProvisionedError(Exception):
pass
# Process-wide app-password storage for the BasicAuth client path.
#
# get_user_client_basic_auth is on the search hot path (Unified Search and the
# /api/v1 viz endpoints call it per request). Creating a fresh
# RefreshTokenStorage and running ``initialize()`` — a full Alembic upgrade in
# a worker thread — on every call is both wasteful and unsafe: concurrent
# upgrades race on Alembic's non-thread-safe module-global EnvironmentContext
# proxy, surfacing as ``KeyError: 'script'``. Cache one initialized instance,
# guarded by a lock so the one-time migration runs exactly once. The lock is
# created lazily inside an async context (anyio primitives must not be built at
# import time — trio compatibility), mirroring vector/qdrant_client.py.
_basic_auth_storage: "RefreshTokenStorage | None" = None
_basic_auth_storage_lock: anyio.Lock | None = None
async def _get_initialized_basic_auth_storage() -> "RefreshTokenStorage":
"""Return the process-wide, already-initialized app-password storage."""
global _basic_auth_storage, _basic_auth_storage_lock
if _basic_auth_storage is not None:
return _basic_auth_storage
# Safe under cooperative scheduling: no await between the None-check and the
# assignment, so two coroutines cannot both create a lock.
if _basic_auth_storage_lock is None:
_basic_auth_storage_lock = anyio.Lock()
async with _basic_auth_storage_lock:
if _basic_auth_storage is None:
storage = RefreshTokenStorage.from_env()
await storage.initialize()
_basic_auth_storage = storage
return _basic_auth_storage
@dataclass
class UserSyncState:
"""State for a single user's scanner task."""
@@ -76,10 +108,11 @@ async def get_user_client_basic_auth(
Raises:
NotProvisionedError: If user has not provisioned an app password
"""
# Get or create storage instance
# Get or create storage instance. Reuse a process-wide initialized instance
# rather than building one (and running an Alembic upgrade) per call — see
# _get_initialized_basic_auth_storage for why (hot path + Alembic race).
if storage is None:
storage = RefreshTokenStorage.from_env()
await storage.initialize()
storage = await _get_initialized_basic_auth_storage()
# Retrieve app password from local storage
app_password = await storage.get_app_password(user_id)