49 lines
1.1 KiB
Python
49 lines
1.1 KiB
Python
"""asyncpg connection pool and DB access helpers."""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
|
|
import asyncpg
|
|
|
|
from app.config import get_settings
|
|
|
|
_pool: asyncpg.Pool | None = None
|
|
|
|
|
|
async def init_db_pool() -> None:
|
|
"""Create global pool (call from FastAPI lifespan)."""
|
|
global _pool
|
|
if _pool is not None:
|
|
return
|
|
s = get_settings()
|
|
_pool = await asyncpg.create_pool(
|
|
host=s.db_host,
|
|
port=s.db_port,
|
|
user=s.db_user,
|
|
password=s.db_password,
|
|
database=s.db_name,
|
|
min_size=1,
|
|
max_size=10,
|
|
)
|
|
|
|
|
|
async def close_db_pool() -> None:
|
|
global _pool
|
|
if _pool is not None:
|
|
await _pool.close()
|
|
_pool = None
|
|
|
|
|
|
def get_pool() -> asyncpg.Pool:
|
|
if _pool is None:
|
|
raise RuntimeError("DB pool not initialized; call init_db_pool() first")
|
|
return _pool
|
|
|
|
|
|
@asynccontextmanager
|
|
async def get_db() -> AsyncIterator[asyncpg.Connection]:
|
|
"""Async context manager yielding a connection from the pool."""
|
|
pool = get_pool()
|
|
async with pool.acquire() as conn:
|
|
yield conn
|