diff --git a/.github/workflows/integration-test.yml b/.github/workflows/integration-test.yml index 4ced849c2..cb404f1fc 100644 --- a/.github/workflows/integration-test.yml +++ b/.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 '*' & diff --git a/backend/open_webui/apps/webui/internal/db.py b/backend/open_webui/apps/webui/internal/db.py index 82dba5031..bcf913e6f 100644 --- a/backend/open_webui/apps/webui/internal/db.py +++ b/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( diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index 5bb5c1238..fbf22d84d 100644 --- a/backend/open_webui/env.py +++ b/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"