From e00f4678df00444a60f0225cdde95d0a35da8263 Mon Sep 17 00:00:00 2001 From: Chris Weaver <25087905+Weves@users.noreply.github.com> Date: Sat, 5 Oct 2024 16:37:48 -0700 Subject: [PATCH] Add option to adjust pool size (#2695) --- backend/danswer/background/update.py | 6 ++++-- backend/danswer/configs/app_configs.py | 6 ++++++ backend/danswer/db/engine.py | 16 ++++++++-------- backend/danswer/main.py | 13 +++++++++---- 4 files changed, 27 insertions(+), 14 deletions(-) diff --git a/backend/danswer/background/update.py b/backend/danswer/background/update.py index c88b349e4..57d05513a 100755 --- a/backend/danswer/background/update.py +++ b/backend/danswer/background/update.py @@ -23,7 +23,7 @@ from danswer.db.connector import fetch_connectors from danswer.db.connector_credential_pair import fetch_connector_credential_pairs from danswer.db.engine import get_db_current_time from danswer.db.engine import get_sqlalchemy_engine -from danswer.db.engine import init_sqlalchemy_engine +from danswer.db.engine import SqlEngine from danswer.db.index_attempt import create_index_attempt from danswer.db.index_attempt import get_index_attempt from danswer.db.index_attempt import get_inprogress_index_attempts @@ -483,7 +483,9 @@ def update_loop( def update__main() -> None: set_is_ee_based_on_env_variable() - init_sqlalchemy_engine(POSTGRES_INDEXER_APP_NAME) + + # initialize the Postgres connection pool + SqlEngine.set_app_name(POSTGRES_INDEXER_APP_NAME) logger.notice("Starting indexing service") update_loop() diff --git a/backend/danswer/configs/app_configs.py b/backend/danswer/configs/app_configs.py index 4857e2aa9..0ac3d6e76 100644 --- a/backend/danswer/configs/app_configs.py +++ b/backend/danswer/configs/app_configs.py @@ -138,6 +138,12 @@ POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "localhost" POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5432" POSTGRES_DB = os.environ.get("POSTGRES_DB") or "postgres" +POSTGRES_API_SERVER_POOL_SIZE = int( + os.environ.get("POSTGRES_API_SERVER_POOL_SIZE") or 40 +) +POSTGRES_API_SERVER_POOL_OVERFLOW = int( + os.environ.get("POSTGRES_API_SERVER_POOL_OVERFLOW") or 10 +) # defaults to False POSTGRES_POOL_PRE_PING = os.environ.get("POSTGRES_POOL_PRE_PING", "").lower() == "true" diff --git a/backend/danswer/db/engine.py b/backend/danswer/db/engine.py index 559c2dec0..dcae5ae1f 100644 --- a/backend/danswer/db/engine.py +++ b/backend/danswer/db/engine.py @@ -26,6 +26,8 @@ from sqlalchemy.orm import sessionmaker from danswer.configs.app_configs import LOG_POSTGRES_CONN_COUNTS from danswer.configs.app_configs import LOG_POSTGRES_LATENCY from danswer.configs.app_configs import MULTI_TENANT +from danswer.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW +from danswer.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE from danswer.configs.app_configs import POSTGRES_DB from danswer.configs.app_configs import POSTGRES_HOST from danswer.configs.app_configs import POSTGRES_PASSWORD @@ -134,8 +136,8 @@ class SqlEngine: # Default parameters for engine creation DEFAULT_ENGINE_KWARGS = { - "pool_size": 40, - "max_overflow": 10, + "pool_size": 20, + "max_overflow": 5, "pool_pre_ping": POSTGRES_POOL_PRE_PING, "pool_recycle": POSTGRES_POOL_RECYCLE, } @@ -201,10 +203,6 @@ def build_connection_string( return f"postgresql+{db_api}://{user}:{password}@{host}:{port}/{db}" -def init_sqlalchemy_engine(app_name: str) -> None: - SqlEngine.set_app_name(app_name) - - def get_sqlalchemy_engine() -> Engine: return SqlEngine.get_engine() @@ -222,8 +220,10 @@ def get_sqlalchemy_async_engine() -> AsyncEngine: "application_name": SqlEngine.get_app_name() + "_async" } }, - pool_size=40, - max_overflow=10, + # async engine is only used by API server, so we can use those values + # here as well + pool_size=POSTGRES_API_SERVER_POOL_SIZE, + max_overflow=POSTGRES_API_SERVER_POOL_OVERFLOW, pool_pre_ping=POSTGRES_POOL_PRE_PING, pool_recycle=POSTGRES_POOL_RECYCLE, ) diff --git a/backend/danswer/main.py b/backend/danswer/main.py index de150cd58..7727f25cc 100644 --- a/backend/danswer/main.py +++ b/backend/danswer/main.py @@ -33,6 +33,8 @@ from danswer.configs.app_configs import DISABLE_INDEX_UPDATE_ON_SWAP from danswer.configs.app_configs import LOG_ENDPOINT_LATENCY from danswer.configs.app_configs import OAUTH_CLIENT_ID from danswer.configs.app_configs import OAUTH_CLIENT_SECRET +from danswer.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW +from danswer.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE from danswer.configs.app_configs import USER_AUTH_SECRET from danswer.configs.app_configs import WEB_DOMAIN from danswer.configs.constants import AuthType @@ -49,8 +51,7 @@ from danswer.db.connector_credential_pair import get_connector_credential_pairs from danswer.db.connector_credential_pair import resync_cc_pair from danswer.db.credentials import create_initial_public_credential from danswer.db.document import check_docs_exist -from danswer.db.engine import get_sqlalchemy_engine -from danswer.db.engine import init_sqlalchemy_engine +from danswer.db.engine import SqlEngine from danswer.db.engine import warm_up_connections from danswer.db.index_attempt import cancel_indexing_attempts_past_model from danswer.db.index_attempt import expire_index_attempts @@ -353,8 +354,12 @@ def setup_vespa( @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncGenerator: - init_sqlalchemy_engine(POSTGRES_WEB_APP_NAME) - engine = get_sqlalchemy_engine() + SqlEngine.set_app_name(POSTGRES_WEB_APP_NAME) + SqlEngine.init_engine( + pool_size=POSTGRES_API_SERVER_POOL_SIZE, + max_overflow=POSTGRES_API_SERVER_POOL_OVERFLOW, + ) + engine = SqlEngine.get_engine() verify_auth = fetch_versioned_implementation( "danswer.auth.users", "verify_auth_setting"