51 lines
1.2 KiB
Python
51 lines
1.2 KiB
Python
"""Database connection and session management"""
|
|
from psycopg import AsyncConnection
|
|
from psycopg_pool import AsyncConnectionPool
|
|
from app.core.config import settings
|
|
from typing import Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Connection pool
|
|
pool: Optional[AsyncConnectionPool] = None
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database connection pool"""
|
|
global pool
|
|
try:
|
|
pool = AsyncConnectionPool(
|
|
conninfo=settings.database_url,
|
|
min_size=1,
|
|
max_size=10,
|
|
open=False,
|
|
)
|
|
await pool.open()
|
|
logger.info("Database connection pool initialized")
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize database pool: {e}")
|
|
raise
|
|
|
|
|
|
async def close_db():
|
|
"""Close database connection pool"""
|
|
global pool
|
|
if pool:
|
|
await pool.close()
|
|
logger.info("Database connection pool closed")
|
|
|
|
|
|
async def get_db() -> AsyncConnection:
|
|
"""Get database connection from pool"""
|
|
if not pool:
|
|
await init_db()
|
|
return await pool.getconn()
|
|
|
|
|
|
async def return_conn(conn: AsyncConnection):
|
|
"""Return connection to pool"""
|
|
if pool:
|
|
await pool.putconn(conn)
|
|
|