Add option to adjust pool size (#2695)

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

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"
 

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,
+        "pool_size": 20,
-        "max_overflow": 10,
+        "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,
+            # async engine is only used by API server, so we can use those values
-            max_overflow=10,
+            # 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,
         )

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 SqlEngine
-from danswer.db.engine import init_sqlalchemy_engine
 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)
+    SqlEngine.set_app_name(POSTGRES_WEB_APP_NAME)
-    engine = get_sqlalchemy_engine()
+    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"