UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2024-11-21 23:57:51 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: db pool config #5935

Browse Source
This commit is contained in:
Sylvere Richard 2024-10-06 16:46:35 +02:00
parent 059ac466e0
commit 9706b76b36
3 changed files with 66 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.github/workflows/integration-test.yml vendored
View File

@ -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 '*' &

26
backend/open_webui/apps/webui/internal/db.py
View File

@ -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(

39
backend/open_webui/env.py
View File

@ -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"