Spaces:
Runtime error
Runtime error
File size: 3,597 Bytes
6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 6b82347 db75ba3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | from contextlib import asynccontextmanager
from typing import AsyncGenerator
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy import text
import os
import dotenv
from database.database_models import Base
from logger.logger import get_logger, log_info, log_error, log_exception
dotenv.load_dotenv()
logger = get_logger(__name__)
class DatabaseManager:
def __init__(self):
self.engine = None
self.async_session_maker = None
self._initialized = False
async def initialize(self):
if self._initialized:
return
database_url = os.getenv("DB_URL")
pool_size = 20
max_overflow = 10
pool_timeout = 30
pool_recycle = 3600
pool_pre_ping = True
self.engine = create_async_engine(
database_url,
echo=False,
pool_size=pool_size,
max_overflow=max_overflow,
pool_timeout=pool_timeout,
pool_recycle=pool_recycle,
pool_pre_ping=pool_pre_ping,
pool_use_lifo=True,
)
self.async_session_maker = async_sessionmaker(
self.engine,
class_=AsyncSession,
expire_on_commit=False,
autocommit=False,
autoflush=False,
)
self._initialized = True
log_info(logger, "Database manager initialized with connection pooling")
async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
if not self._initialized:
await self.initialize()
async with self.async_session_maker() as session:
try:
yield session
except Exception as e:
await session.rollback()
log_exception(logger, e)
raise
finally:
await session.close()
@asynccontextmanager
async def connect(self):
if not self._initialized:
await self.initialize()
async with self.async_session_maker() as session:
try:
yield session
finally:
await session.close()
async def get_pool_stats(self) -> dict:
if not self._initialized or not self.engine:
return {"error": "Database not initialized"}
pool = self.engine.pool
return {
"size": pool.size(),
"checked_in": pool.checkedin(),
"overflow": pool.overflow(),
"total": pool.size() + pool.overflow(),
"checked_out": pool.checkedout(),
"initialized": self._initialized,
}
async def health_check(self) -> bool:
try:
async with self.connect() as session:
await session.execute(text("SELECT 1"))
return True
except Exception as e:
log_error(logger, f"Health check failed: {str(e)}")
return False
async def create_tables(self):
await self.initialize()
async with self.engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
async def close_all(self):
if self.engine:
await self.engine.dispose()
self._initialized = False
log_info(logger, "All database connections closed")
db_manager = DatabaseManager()
|