Fix: Thread-safe DB connection manager with pooling (#212)#252
Fix: Thread-safe DB connection manager with pooling (#212)#252muhammadtihame wants to merge 2 commits into
Conversation
📝 WalkthroughWalkthroughAdds async database support via SQLAlchemy/asyncpg: new optional Changes
Sequence Diagram(s)sequenceDiagram
participant Client as FastAPI Route
participant GetDb as get_db()
participant Maker as async_session_maker
participant Pool as Connection Pool
participant Session as AsyncSession
Client->>GetDb: request session
alt Engine initialized
GetDb->>Maker: create session
Maker->>Pool: acquire connection
Pool-->>Maker: provide connection
Maker-->>Session: return AsyncSession
GetDb-->>Client: yield session
Client->>Session: execute queries
alt Success
Client->>GetDb: return normally
GetDb->>Session: close session
Session->>Pool: release connection
else Exception
Client->>GetDb: exception raised
GetDb->>Session: rollback
GetDb->>Session: close session
Session->>Pool: release connection
end
else Engine not initialized
GetDb-->>Client: raise RuntimeError
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@backend/app/database/core.py`:
- Around line 46-57: Remove the redundant explicit close in the async session
context manager: when using the async with async_session_maker() as session
block you should not call await session.close() in the finally; delete that call
so the context manager handles cleanup. Also replace logger.error(...) with
logger.exception(...) in the except block to include the traceback, while
preserving await session.rollback() and re-raising the exception; reference
async_session_maker, session, logger.error -> logger.exception, session.rollback
to locate the changes.
In `@tests/test_db_pool.py`:
- Around line 7-9: The test currently patches app.core.config.settings at module
scope but imports engine and get_db after the patch context, which is fragile
because engine is created at import time and module caching prevents
reinitialization; update tests/test_db_pool.py to apply the patch while
importing/reloading the database module so engine is created with the mocked
configuration—use a pytest fixture that applies patch.dict or patch to set
DATABASE_URL (or patch app.core.config.settings), then
importlib.reload(app.database.core) inside that fixture and yield db_core.engine
and db_core.get_db (or move the import of engine/get_db inside the with
patch(...) block) so engine initialization sees the mocked value.
- Around line 44-45: The async context manager mock is incomplete: instead of
assigning __aexit__.return_value = None make both __aenter__ and __aexit__
awaitable (e.g., replace the current assignments on mock_session with AsyncMock
usage: set mock_session.__aenter__ = AsyncMock(return_value=mock_session) and
mock_session.__aexit__ = AsyncMock(return_value=None)), and apply the same
change to the other async context-manager mock later in the file so both
enter/exit are proper awaitables.
🧹 Nitpick comments (1)
tests/test_db_pool.py (1)
49-52: Unused loop variablesession— rename to_session.Per the static analysis hint, the loop variable is unused within the body. Renaming to
_sessionsignals intentional non-use.- async for session in get_db(): + async for _session in get_db():The same applies to line 80.
|
Fixes #212 |
|
Addressed CodeRabbit feedback:
|
Shevilll
left a comment
Shevilll
left a comment
There was a problem hiding this comment.
Peer Review: Fix: Thread-safe DB connection manager with pooling
Hello @muhammadtihame! Thank you for this massive and extremely important improvement to the database layer. Transitioning the backend to use asynchronous SQLAlchemy connection pooling via asyncpg is a huge win for concurrent performance and scalability, especially when handling high API traffic compared to the old direct postgrest calls.
I've reviewed the code in detail and have a few constructive suggestions to share:
1. 🧹 PEP 8 Style and Readability
In backend/app/database/core.py, both engine and async_session_maker are initialized using trailing ternary if/else statements wrapped around long, multiline function calls:
engine = create_async_engine(
DATABASE_URL,
echo=False,
pool_size=20,
...
) if DATABASE_URL else NoneWhile syntactically valid in Python, putting a multiline block inside a ternary condition makes the flow harder to parse and violates PEP 8's readability guidelines. It's much cleaner to initialize these using standard if blocks:
engine = None
async_session_maker = None
if DATABASE_URL:
engine = create_async_engine(
DATABASE_URL,
echo=False,
pool_size=20, # Maintain 20 open connections
max_overflow=10, # Allow 10 extra during spikes
pool_timeout=30, # Wait 30s for a connection before raising timeout
pool_pre_ping=True, # Check connection health before handing it out
)
async_session_maker = async_sessionmaker(
engine,
class_=AsyncSession,
expire_on_commit=False,
autocommit=False,
autoflush=False,
)This makes the conditional initialization of both the engine and session factory explicit, clean, and extremely easy to read.
2. 🛡️ Robust Async Session Cleanup
In the get_db generator:
async with async_session_maker() as session:
try:
yield session
except Exception:
logger.exception("Database session error")
await session.rollback()
raiseThis is a robust and safe pattern!
Just as a small technical detail to note: if a request is cancelled or the connection is severed, FastAPI throws a GeneratorExit (which inherits from BaseException rather than Exception) into the generator at the yield statement.
Because except Exception only catches standard exceptions, GeneratorExit will propagate past this block. However, the outer async with block will still catch it and execute its __aexit__ method, which automatically closes the SQLAlchemy session (performing a safe rollback of any active transactions). Thus, your implementation is completely safe even during task cancellations!
3. 🧪 Commendation on Tests
The test suite you added in tests/test_db_pool.py is absolutely phenomenal! Testing database connection pools and concurrency with multiple async tasks is historically very difficult to mock correctly, but your mock set up (simulating 50 parallel requests and verifying rollback and cleanup counts) is exceptionally clean, robust, and elegant.
Outstanding work on this! Once the stylistic ternary block is cleaned up, this is ready to go.
Shevilll
left a comment
Shevilll
left a comment
There was a problem hiding this comment.
Peer Review: Thread-safe DB connection manager with pooling
Thank you for putting together this implementation of the database connection pooling using SQLAlchemy and asyncpg. It is a solid foundation, but there are a few architectural, thread-safety, and operational issues that need to be addressed before this is ready for production.
1. Severe Log Pollution via Global Exception Catching in get_db()
In backend/app/database/core.py, the dependency catches all exceptions that propagate through the yield:
async def get_db() -> AsyncGenerator[AsyncSession, None]:
...
async with async_session_maker() as session:
try:
yield session
except Exception:
logger.exception("Database session error")
await session.rollback()
raiseIssue: Any exception raised inside the FastAPI endpoint or downstream handlers (including standard validation errors like RequestValidationError / HTTP 422, or standard not-found exceptions like HTTPException / HTTP 404) bubbles back through the generator. By catching all Exception types and logging them with logger.exception("Database session error"), standard client-side errors will produce scary server-side traceback logs. This will pollute production logging systems and trigger false-alarm alerts.
Fix: Limit the logging to actual database-level exceptions (e.g., sqlalchemy.exc.SQLAlchemyError), or perform the rollback silently (or at DEBUG/INFO level) without logging general application exceptions as database failures:
try:
yield session
except sqlalchemy.exc.SQLAlchemyError as e:
logger.exception("Database session database-level error")
await session.rollback()
raise
except Exception:
# Rollback for safety, but do not log as a DB system error
await session.rollback()
raise2. Standard PostgreSQL URIs Will Crash Async Engine Initialization
create_async_engine is invoked with DATABASE_URL directly:
engine = create_async_engine(
DATABASE_URL,
...
) if DATABASE_URL else NoneIssue: Standard connection strings (e.g., from Supabase or PostgreSQL environments) start with postgresql:// or postgres://. SQLAlchemy's create_async_engine requires an async dialect (like postgresql+asyncpg://). If a standard URL is passed, it will fail to initialize with InvalidRequestError: The asyncpg dialect is not installed or similar.
Fix: Implement automatic scheme conversion in the settings layer or right before creating the engine:
if DATABASE_URL:
if DATABASE_URL.startswith("postgresql://"):
DATABASE_URL = DATABASE_URL.replace("postgresql://", "postgresql+asyncpg://", 1)
elif DATABASE_URL.startswith("postgres://"):
DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql+asyncpg://", 1)3. Module-Level Engine Creation Anti-Pattern
Issue: Creating the engine and sessionmaker at the module-import level (backend/app/database/core.py) means any configuration or parse error will crash the application during import. This makes running CLI scripts, setup/migration scripts, or testing suites that do not use the database highly brittle.
Fix: Wrap the engine and sessionmaker in a lazy-initialization block, or initialize them as part of a FastAPI lifespan event handler.
4. Brittle, Stateful Unit Testing via importlib.reload
In tests/test_db_pool.py:
@pytest.fixture
def mock_db_module():
with patch("app.core.config.settings") as mock_settings:
mock_settings.database_url = "postgresql+asyncpg://user:password@localhost:5432/testdb"
import app.database.core
importlib.reload(app.database.core)
yield app.database.coreIssue: Using importlib.reload during tests is a major anti-pattern. It can lead to state leakages, where other test modules that imported app.database.core retain references to the old engine or get confused by partial reloads.
Fix: Cleanly decouple settings injection from the connection manager, allowing settings to be passed or resolved dynamically rather than relying on module reloads.
2 participants
Summary
This PR implements thread-safe database connection management with connection pooling and proper resource cleanup.
Changes
backend/app/database/core.py(async SQLAlchemy engine + session manager)tests/test_db_pool.pyto verify concurrency and rollback behaviordocs/DATABASE_CONNECTION.mddocumentationbackend/app/core/config/settings.pyfor database configurationpyproject.tomlandbackend/requirements.txtto include new dependenciesVerification
Run the following tests: