78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
"""
|
|
Base database configuration for all microservices
|
|
"""
|
|
|
|
import os
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
|
from sqlalchemy.orm import sessionmaker, declarative_base
|
|
from sqlalchemy.pool import StaticPool
|
|
from contextlib import asynccontextmanager
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
Base = declarative_base()
|
|
|
|
class DatabaseManager:
|
|
"""Database manager for microservices"""
|
|
|
|
def __init__(self, database_url: str):
|
|
self.database_url = database_url
|
|
self.async_engine = create_async_engine(
|
|
database_url,
|
|
echo=False,
|
|
pool_pre_ping=True,
|
|
pool_recycle=300,
|
|
pool_size=20,
|
|
max_overflow=30
|
|
)
|
|
|
|
self.async_session_local = sessionmaker(
|
|
self.async_engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False
|
|
)
|
|
|
|
async def get_db(self):
|
|
"""Get database session for request handlers"""
|
|
async with self.async_session_local() as session:
|
|
try:
|
|
yield session
|
|
except Exception as e:
|
|
logger.error(f"Database session error: {e}")
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
@asynccontextmanager
|
|
async def get_background_session(self):
|
|
"""
|
|
✅ NEW: Get database session for background tasks
|
|
|
|
Usage:
|
|
async with database_manager.get_background_session() as session:
|
|
# Your background task code here
|
|
await session.commit()
|
|
"""
|
|
async with self.async_session_local() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception as e:
|
|
await session.rollback()
|
|
logger.error(f"Background task database error: {e}")
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
async def create_tables(self):
|
|
"""Create database tables"""
|
|
async with self.async_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
async def drop_tables(self):
|
|
"""Drop database tables"""
|
|
async with self.async_engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.drop_all) |