feat: db pool config #5935

.github/workflows/integration-test.yml

@@ -182,6 +182,9 @@ jobs:
           WEBUI_SECRET_KEY: secret-key
           GLOBAL_LOG_LEVEL: debug
           DATABASE_URL: postgresql://postgres:postgres@localhost:5432/postgres
+          DATABASE_POOL_SIZE: 10
+          DATABASE_POOL_MAX_OVERFLOW: 10
+          DATABASE_POOL_TIMEOUT: 30
         run: |
           cd backend
           uvicorn open_webui.main:app --port "8081" --forwarded-allow-ips '*' &

backend/open_webui/apps/webui/internal/db.py

@@ -4,11 +4,20 @@ from contextlib import contextmanager
 from typing import Any, Optional
 
 from open_webui.apps.webui.internal.wrappers import register_connection
-from open_webui.env import OPEN_WEBUI_DIR, DATABASE_URL, SRC_LOG_LEVELS
+from open_webui.env import (
+    OPEN_WEBUI_DIR,
+    DATABASE_URL,
+    SRC_LOG_LEVELS,
+    DATABASE_POOL_MAX_OVERFLOW,
+    DATABASE_POOL_RECYCLE,
+    DATABASE_POOL_SIZE,
+    DATABASE_POOL_TIMEOUT,
+)
 from peewee_migrate import Router
 from sqlalchemy import Dialect, create_engine, types
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import scoped_session, sessionmaker
+from sqlalchemy.pool import QueuePool, NullPool
 from sqlalchemy.sql.type_api import _T
 from typing_extensions import Self
 
@@ -71,7 +80,20 @@ if "sqlite" in SQLALCHEMY_DATABASE_URL:
         SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
     )
 else:
-    engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
+    if DATABASE_POOL_SIZE > 0:
+        engine = create_engine(
+            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:
+        engine = create_engine(
+            SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, poolclass=NullPool
+        )
 
 
 SessionLocal = sessionmaker(

backend/open_webui/env.py

@@ -258,6 +258,45 @@ DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite:///{DATA_DIR}/webui.db")
 if "postgres://" in DATABASE_URL:
     DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://")
 
+DATABASE_POOL_SIZE = os.environ.get("DATABASE_POOL_SIZE", 0)
+
+if DATABASE_POOL_SIZE == "":
+    DATABASE_POOL_SIZE = 0
+else:
+    try:
+        DATABASE_POOL_SIZE = int(DATABASE_POOL_SIZE)
+    except Exception:
+        DATABASE_POOL_SIZE = 0
+
+DATABASE_POOL_MAX_OVERFLOW = os.environ.get("DATABASE_POOL_MAX_OVERFLOW", 0)
+
+if DATABASE_POOL_MAX_OVERFLOW == "":
+    DATABASE_POOL_MAX_OVERFLOW = 0
+else:
+    try:
+        DATABASE_POOL_MAX_OVERFLOW = int(DATABASE_POOL_MAX_OVERFLOW)
+    except Exception:
+        DATABASE_POOL_MAX_OVERFLOW = 0
+
+DATABASE_POOL_TIMEOUT = os.environ.get("DATABASE_POOL_TIMEOUT", 30)
+
+if DATABASE_POOL_TIMEOUT == "":
+    DATABASE_POOL_TIMEOUT = 30
+else:
+    try:
+        DATABASE_POOL_TIMEOUT = int(DATABASE_POOL_TIMEOUT)
+    except Exception:
+        DATABASE_POOL_TIMEOUT = 30
+
+DATABASE_POOL_RECYCLE = os.environ.get("DATABASE_POOL_RECYCLE", 3600)
+
+if DATABASE_POOL_RECYCLE == "":
+    DATABASE_POOL_RECYCLE = 3600
+else:
+    try:
+        DATABASE_POOL_RECYCLE = int(DATABASE_POOL_RECYCLE)
+    except Exception:
+        DATABASE_POOL_RECYCLE = 3600
 
 RESET_CONFIG_ON_START = (
     os.environ.get("RESET_CONFIG_ON_START", "False").lower() == "true"