From e8638227fd94114acbb03fa4a365cd847c85b693 Mon Sep 17 00:00:00 2001
From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Date: Tue, 27 May 2025 17:27:53 -0400
Subject: [PATCH] fix (config): Configurable Pool Class

---
 src/config.py |  1 +
 src/db.py     | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/src/config.py b/src/config.py
index 33161d4a..1828b10c 100644
--- a/src/config.py
+++ b/src/config.py
@@ -113,6 +113,7 @@ class DBSettings(TomlSettings):
         "postgresql+psycopg://postgres:postgres@localhost:5432/postgres"
     )
     SCHEMA: str = "public"
+    POOL_CLASS: str = "default"
     POOL_PRE_PING: bool = True
     POOL_SIZE: int = 10
     MAX_OVERFLOW: int = 20
diff --git a/src/db.py b/src/db.py
index 317bb011..32d16170 100644
--- a/src/db.py
+++ b/src/db.py
@@ -4,6 +4,7 @@ from typing import Optional
 from sqlalchemy import MetaData, create_engine, text
 from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
 from sqlalchemy.orm import declarative_base
+from sqlalchemy.pool import NullPool, AsyncAdaptedQueuePool
 
 from src.config import settings
 
@@ -15,6 +16,24 @@ request_context: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
     default=None,
 )
 
+engine_kwargs = {}
+
+pool_class = AsyncAdaptedQueuePool
+if settings.DB.POOL_CLASS == "null":
+    engine_kwargs["poolclass"] = NullPool
+else:
+    # Only add pool-related kwargs for pooled connections
+    engine_kwargs.update(
+        {
+            "pool_pre_ping": settings.DB.POOL_PRE_PING,
+            "pool_size": settings.DB.POOL_SIZE,
+            "max_overflow": settings.DB.MAX_OVERFLOW,
+            "pool_timeout": settings.DB.POOL_TIMEOUT,
+            "pool_recycle": settings.DB.POOL_RECYCLE,
+            "pool_use_lifo": settings.DB.POOL_USE_LIFO,
+        }
+    )
+
 engine = create_async_engine(
     settings.DB.CONNECTION_URI,
     connect_args=connect_args,
@@ -25,6 +44,7 @@ engine = create_async_engine(
     pool_timeout=settings.DB.POOL_TIMEOUT,
     pool_recycle=settings.DB.POOL_RECYCLE,
     pool_use_lifo=settings.DB.POOL_USE_LIFO,
+    poolclass=pool_class,
 )
 
 SessionLocal = async_sessionmaker(