c0371d85ce
- Removed direct pool checks and replaced them with a centralized database initialization method in `init_db`. - Updated API endpoints in `admin.py`, `collection.py`, `pins.py`, and `watched.py` to ensure the database connection pool is initialized before usage. - Enhanced error handling to raise HTTP exceptions if the database is unavailable. - Improved the `init_db` function in `database.py` to prevent multiple simultaneous initializations using an asyncio lock.
75 lines
1.9 KiB
Python
75 lines
1.9 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
|
|
import asyncio
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Connection pool
|
|
pool: Optional[AsyncConnectionPool] = None
|
|
_init_lock = asyncio.Lock()
|
|
_initializing = False
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database connection pool"""
|
|
global pool, _initializing
|
|
|
|
# If already initialized, return
|
|
if pool is not None:
|
|
return
|
|
|
|
# Use lock to prevent multiple simultaneous initializations
|
|
async with _init_lock:
|
|
# Double-check after acquiring lock
|
|
if pool is not None:
|
|
return
|
|
|
|
if _initializing:
|
|
# Wait for other initialization to complete
|
|
while _initializing:
|
|
await asyncio.sleep(0.1)
|
|
return
|
|
|
|
_initializing = True
|
|
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}")
|
|
pool = None
|
|
raise
|
|
finally:
|
|
_initializing = False
|
|
|
|
|
|
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)
|
|
|