fix: release database connections immediately after auth instead of holding during LLM calls (#20545)
fix: release database connections immediately after auth instead of holding during LLM calls
Authentication was using Depends(get_session) which holds a database connection
for the entire request lifecycle. For chat completions, this meant connections
were held for 30-60 seconds while waiting for LLM responses, despite only needing
the connection for ~50ms of actual database work.
With a default pool of 15 connections, this limited concurrent chat users to ~15
before pool exhaustion and timeout errors:
sqlalchemy.exc.TimeoutError: QueuePool limit of size 5 overflow 10 reached,
connection timed out, timeout 30.00
The fix removes Depends(get_session) from get_current_user. Each database
operation now manages its own short-lived session internally:
BEFORE: One session held for entire request
──────────────────────────────────────────────────
│ auth │ queries │ LLM wait (30s) │ save │
│ CONNECTION HELD ENTIRE TIME │
──────────────────────────────────────────────────
AFTER: Short-lived sessions, released immediately
┌──────┐ ┌───────┐ ┌──────┐
│ auth │ │ query │ LLM (30s) │ save │
│ 10ms │ │ 20ms │ NO CONNECTION │ 20ms │
└──────┘ └───────┘ └──────┘
This is safe because:
- User model has no lazy-loaded relationships (all simple columns)
- Pydantic conversion (UserModel.model_validate) happens while session is open
- Returned object is pure Pydantic with no SQLAlchemy ties
Combined with the telemetry efficiency fix, this resolves connection pool
exhaustion for high-concurrency deployments, particularly on network-attached
databases like AWS Aurora where connection hold time is more impactful.
This commit is contained in:
@@ -44,8 +44,6 @@ from open_webui.env import (
|
||||
|
||||
from fastapi import BackgroundTasks, Depends, HTTPException, Request, Response, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
from sqlalchemy.orm import Session
|
||||
from open_webui.internal.db import get_session
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
@@ -279,7 +277,10 @@ async def get_current_user(
|
||||
response: Response,
|
||||
background_tasks: BackgroundTasks,
|
||||
auth_token: HTTPAuthorizationCredentials = Depends(bearer_security),
|
||||
db: Session = Depends(get_session),
|
||||
# NOTE: We intentionally do NOT use Depends(get_session) here.
|
||||
# Sessions are managed internally with short-lived context managers.
|
||||
# This ensures connections are released immediately after auth queries,
|
||||
# not held for the entire request duration (e.g., during 30+ second LLM calls).
|
||||
):
|
||||
token = None
|
||||
|
||||
@@ -294,7 +295,7 @@ async def get_current_user(
|
||||
|
||||
# auth by api key
|
||||
if token.startswith("sk-"):
|
||||
user = get_current_user_by_api_key(request, token, db=db)
|
||||
user = get_current_user_by_api_key(request, token)
|
||||
|
||||
# Add user info to current span
|
||||
current_span = trace.get_current_span()
|
||||
@@ -323,7 +324,7 @@ async def get_current_user(
|
||||
detail="Invalid token",
|
||||
)
|
||||
|
||||
user = Users.get_user_by_id(data["id"], db=db)
|
||||
user = Users.get_user_by_id(data["id"])
|
||||
if user is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
@@ -373,8 +374,9 @@ async def get_current_user(
|
||||
raise e
|
||||
|
||||
|
||||
def get_current_user_by_api_key(request, api_key: str, db: Session = None):
|
||||
user = Users.get_user_by_api_key(api_key, db=db)
|
||||
def get_current_user_by_api_key(request, api_key: str):
|
||||
# Each function call manages its own short-lived session internally
|
||||
user = Users.get_user_by_api_key(api_key)
|
||||
|
||||
if user is None:
|
||||
raise HTTPException(
|
||||
@@ -402,7 +404,7 @@ def get_current_user_by_api_key(request, api_key: str, db: Session = None):
|
||||
current_span.set_attribute("client.user.role", user.role)
|
||||
current_span.set_attribute("client.auth.type", "api_key")
|
||||
|
||||
Users.update_last_active_by_id(user.id, db=db)
|
||||
Users.update_last_active_by_id(user.id)
|
||||
return user
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user