refac: async db

backend/open_webui/internal/db.py

@@ -1,7 +1,7 @@
 import os
 import json
 import logging
-from contextlib import contextmanager
+from contextlib import asynccontextmanager, contextmanager
 from typing import Any, Optional
 
 from open_webui.internal.wrappers import register_connection
@@ -19,6 +19,7 @@ from open_webui.env import (
 )
 from peewee_migrate import Router
 from sqlalchemy import Dialect, create_engine, MetaData, event, types
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import scoped_session, sessionmaker, Session
 from sqlalchemy.pool import QueuePool, NullPool
@@ -81,6 +82,32 @@ if ENABLE_DB_MIGRATIONS:
 
 SQLALCHEMY_DATABASE_URL = DATABASE_URL
 
+
+def _make_async_url(url: str) -> str:
+    """Convert a sync database URL to its async driver equivalent."""
+    if url.startswith('sqlite+sqlcipher://'):
+        # SQLCipher has no async driver — not supported for async
+        raise ValueError(
+            'sqlite+sqlcipher:// URLs are not supported with async engine. '
+            'Use standard sqlite:// or postgresql:// instead.'
+        )
+    if url.startswith('sqlite:///') or url.startswith('sqlite://'):
+        return url.replace('sqlite://', 'sqlite+aiosqlite://', 1)
+    if url.startswith('postgresql+psycopg2://'):
+        return url.replace('postgresql+psycopg2://', 'postgresql+asyncpg://', 1)
+    if url.startswith('postgresql://'):
+        return url.replace('postgresql://', 'postgresql+asyncpg://', 1)
+    if url.startswith('postgres://'):
+        return url.replace('postgres://', 'postgresql+asyncpg://', 1)
+    # For other dialects, return as-is and let SQLAlchemy handle it
+    return url
+
+
+# ============================================================
+# SYNC ENGINE (used only for: startup migrations, config loading,
+#              Alembic, peewee migration, health checks)
+# ============================================================
+
 # Handle SQLCipher URLs
 if SQLALCHEMY_DATABASE_URL.startswith('sqlite+sqlcipher://'):
     database_password = os.environ.get('DATABASE_PASSWORD')
@@ -155,6 +182,7 @@ else:
         engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
 
 
+# Sync session — used ONLY for startup config loading (config.py runs at import time)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, expire_on_commit=False)
 metadata_obj = MetaData(schema=DATABASE_SCHEMA)
 Base = declarative_base(metadata=metadata_obj)
@@ -162,6 +190,7 @@ ScopedSession = scoped_session(SessionLocal)
 
 
 def get_session():
+    """Sync session generator — used ONLY for startup/config operations."""
     db = SessionLocal()
     try:
         yield db
@@ -172,10 +201,82 @@ def get_session():
 get_db = contextmanager(get_session)
 
 
-@contextmanager
-def get_db_context(db: Optional[Session] = None):
-    if isinstance(db, Session) and DATABASE_ENABLE_SESSION_SHARING:
+# ============================================================
+# ASYNC ENGINE (used for ALL runtime database operations)
+# ============================================================
+
+ASYNC_SQLALCHEMY_DATABASE_URL = _make_async_url(SQLALCHEMY_DATABASE_URL)
+
+if 'sqlite' in ASYNC_SQLALCHEMY_DATABASE_URL:
+    async_engine = create_async_engine(
+        ASYNC_SQLALCHEMY_DATABASE_URL,
+        connect_args={'check_same_thread': False},
+    )
+
+    if DATABASE_ENABLE_SQLITE_WAL:
+        @event.listens_for(async_engine.sync_engine, 'connect')
+        def _set_sqlite_wal(dbapi_connection, connection_record):
+            cursor = dbapi_connection.cursor()
+            cursor.execute('PRAGMA journal_mode=WAL')
+            cursor.close()
+else:
+    if isinstance(DATABASE_POOL_SIZE, int):
+        if DATABASE_POOL_SIZE > 0:
+            async_engine = create_async_engine(
+                ASYNC_SQLALCHEMY_DATABASE_URL,
+                pool_size=DATABASE_POOL_SIZE,
+                max_overflow=DATABASE_POOL_MAX_OVERFLOW,
+                pool_timeout=DATABASE_POOL_TIMEOUT,
+                pool_recycle=DATABASE_POOL_RECYCLE,
+                pool_pre_ping=True,
+                poolclass=QueuePool,
+            )
+        else:
+            async_engine = create_async_engine(
+                ASYNC_SQLALCHEMY_DATABASE_URL,
+                pool_pre_ping=True,
+                poolclass=NullPool,
+            )
+    else:
+        async_engine = create_async_engine(
+            ASYNC_SQLALCHEMY_DATABASE_URL,
+            pool_pre_ping=True,
+        )
+
+
+AsyncSessionLocal = async_sessionmaker(
+    bind=async_engine,
+    class_=AsyncSession,
+    autocommit=False,
+    autoflush=False,
+    expire_on_commit=False,
+)
+
+
+async def get_async_session():
+    """Async session generator for FastAPI Depends()."""
+    async with AsyncSessionLocal() as db:
+        try:
+            yield db
+        finally:
+            await db.close()
+
+
+@asynccontextmanager
+async def get_async_db():
+    """Async context manager for use outside of FastAPI dependency injection."""
+    async with AsyncSessionLocal() as db:
+        try:
+            yield db
+        finally:
+            await db.close()
+
+
+@asynccontextmanager
+async def get_async_db_context(db: Optional[AsyncSession] = None):
+    """Async context manager that reuses an existing session if provided and session sharing is enabled."""
+    if isinstance(db, AsyncSession) and DATABASE_ENABLE_SESSION_SHARING:
         yield db
     else:
-        with get_db() as session:
+        async with get_async_db() as session:
             yield session