REFACTOR - Database logic
This commit is contained in:
Urtzi Alfaro
@@ -1,46 +1,196 @@
|
||||
"""Database configuration for data service"""
|
||||
"""
|
||||
Database configuration for data service
|
||||
Uses shared database infrastructure for consistency
|
||||
"""
|
||||
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||||
from sqlalchemy.orm import declarative_base
|
||||
import structlog
|
||||
from typing import AsyncGenerator
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import text
|
||||
|
||||
from shared.database.base import DatabaseManager, Base
|
||||
from app.core.config import settings
|
||||
|
||||
logger = structlog.get_logger()
|
||||
|
||||
# Create async engine
|
||||
engine = create_async_engine(
|
||||
settings.DATABASE_URL,
|
||||
echo=False,
|
||||
pool_pre_ping=True,
|
||||
pool_size=10,
|
||||
max_overflow=20
|
||||
# Initialize database manager using shared infrastructure
|
||||
database_manager = DatabaseManager(
|
||||
database_url=settings.DATABASE_URL,
|
||||
service_name="data",
|
||||
pool_size=15,
|
||||
max_overflow=25,
|
||||
echo=settings.DEBUG if hasattr(settings, 'DEBUG') else False
|
||||
)
|
||||
|
||||
# Create async session factory
|
||||
AsyncSessionLocal = async_sessionmaker(
|
||||
engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False
|
||||
)
|
||||
# Alias for convenience - matches the existing interface
|
||||
get_db = database_manager.get_db
|
||||
|
||||
# Base class for models
|
||||
Base = declarative_base()
|
||||
# Use the shared background session method
|
||||
get_background_db_session = database_manager.get_background_session
|
||||
|
||||
async def get_db() -> AsyncSession:
|
||||
"""Get database session"""
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
await session.close()
|
||||
async def get_db_health() -> bool:
|
||||
"""Health check function for database connectivity"""
|
||||
try:
|
||||
async with database_manager.async_engine.begin() as conn:
|
||||
await conn.execute(text("SELECT 1"))
|
||||
logger.debug("Database health check passed")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Database health check failed", error=str(e))
|
||||
return False
|
||||
|
||||
async def init_db():
|
||||
"""Initialize database tables"""
|
||||
"""Initialize database tables using shared infrastructure"""
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
logger.info("Database initialized successfully")
|
||||
logger.info("Initializing data service database")
|
||||
|
||||
# Import models to ensure they're registered
|
||||
from app.models.sales import SalesData
|
||||
from app.models.traffic import TrafficData
|
||||
from app.models.weather import WeatherData
|
||||
|
||||
# Create tables using shared infrastructure
|
||||
await database_manager.create_tables()
|
||||
|
||||
logger.info("Data service database initialized successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize database", error=str(e))
|
||||
raise
|
||||
logger.error("Failed to initialize data service database", error=str(e))
|
||||
raise
|
||||
|
||||
# Data service specific database utilities
|
||||
class DataDatabaseUtils:
|
||||
"""Data service specific database utilities"""
|
||||
|
||||
@staticmethod
|
||||
async def cleanup_old_sales_data(days_old: int = 730):
|
||||
"""Clean up old sales data (default 2 years)"""
|
||||
try:
|
||||
async with database_manager.get_background_session() as session:
|
||||
if settings.DATABASE_URL.startswith("sqlite"):
|
||||
query = text(
|
||||
"DELETE FROM sales_data "
|
||||
"WHERE created_at < datetime('now', :days_param)"
|
||||
)
|
||||
params = {"days_param": f"-{days_old} days"}
|
||||
else:
|
||||
query = text(
|
||||
"DELETE FROM sales_data "
|
||||
"WHERE created_at < NOW() - INTERVAL :days_param"
|
||||
)
|
||||
params = {"days_param": f"{days_old} days"}
|
||||
|
||||
result = await session.execute(query, params)
|
||||
deleted_count = result.rowcount
|
||||
|
||||
logger.info("Cleaned up old sales data",
|
||||
deleted_count=deleted_count,
|
||||
days_old=days_old)
|
||||
|
||||
return deleted_count
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to cleanup old sales data", error=str(e))
|
||||
return 0
|
||||
|
||||
@staticmethod
|
||||
async def get_data_statistics(tenant_id: str = None) -> dict:
|
||||
"""Get data service statistics"""
|
||||
try:
|
||||
async with database_manager.get_background_session() as session:
|
||||
# Get sales data statistics
|
||||
if tenant_id:
|
||||
sales_query = text(
|
||||
"SELECT COUNT(*) as count "
|
||||
"FROM sales_data "
|
||||
"WHERE tenant_id = :tenant_id"
|
||||
)
|
||||
params = {"tenant_id": tenant_id}
|
||||
else:
|
||||
sales_query = text("SELECT COUNT(*) as count FROM sales_data")
|
||||
params = {}
|
||||
|
||||
sales_result = await session.execute(sales_query, params)
|
||||
sales_count = sales_result.scalar() or 0
|
||||
|
||||
# Get traffic data statistics (if exists)
|
||||
try:
|
||||
traffic_query = text("SELECT COUNT(*) as count FROM traffic_data")
|
||||
if tenant_id:
|
||||
# Traffic data might not have tenant_id, check table structure
|
||||
pass
|
||||
|
||||
traffic_result = await session.execute(traffic_query)
|
||||
traffic_count = traffic_result.scalar() or 0
|
||||
except:
|
||||
traffic_count = 0
|
||||
|
||||
# Get weather data statistics (if exists)
|
||||
try:
|
||||
weather_query = text("SELECT COUNT(*) as count FROM weather_data")
|
||||
weather_result = await session.execute(weather_query)
|
||||
weather_count = weather_result.scalar() or 0
|
||||
except:
|
||||
weather_count = 0
|
||||
|
||||
return {
|
||||
"tenant_id": tenant_id,
|
||||
"sales_records": sales_count,
|
||||
"traffic_records": traffic_count,
|
||||
"weather_records": weather_count,
|
||||
"total_records": sales_count + traffic_count + weather_count
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to get data statistics", error=str(e))
|
||||
return {
|
||||
"tenant_id": tenant_id,
|
||||
"sales_records": 0,
|
||||
"traffic_records": 0,
|
||||
"weather_records": 0,
|
||||
"total_records": 0
|
||||
}
|
||||
|
||||
# Enhanced database session dependency with better error handling
|
||||
async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
"""Enhanced database session dependency with better logging and error handling"""
|
||||
async with database_manager.async_session_local() as session:
|
||||
try:
|
||||
logger.debug("Database session created")
|
||||
yield session
|
||||
except Exception as e:
|
||||
logger.error("Database session error", error=str(e), exc_info=True)
|
||||
await session.rollback()
|
||||
raise
|
||||
finally:
|
||||
await session.close()
|
||||
logger.debug("Database session closed")
|
||||
|
||||
# Database cleanup for data service
|
||||
async def cleanup_data_database():
|
||||
"""Cleanup database connections for data service"""
|
||||
try:
|
||||
logger.info("Cleaning up data service database connections")
|
||||
|
||||
# Close engine connections
|
||||
if hasattr(database_manager, 'async_engine') and database_manager.async_engine:
|
||||
await database_manager.async_engine.dispose()
|
||||
|
||||
logger.info("Data service database cleanup completed")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to cleanup data service database", error=str(e))
|
||||
|
||||
# Export the commonly used items to maintain compatibility
|
||||
__all__ = [
|
||||
'Base',
|
||||
'database_manager',
|
||||
'get_db',
|
||||
'get_background_db_session',
|
||||
'get_db_session',
|
||||
'get_db_health',
|
||||
'DataDatabaseUtils',
|
||||
'init_db',
|
||||
'cleanup_data_database'
|
||||
]
|
||||
Reference in New Issue
Block a user