""" Database configuration for data service Uses shared database infrastructure for consistency """ 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() # 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 ) # Alias for convenience - matches the existing interface get_db = database_manager.get_db # Use the shared background session method get_background_db_session = database_manager.get_background_session 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 using shared infrastructure""" try: 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 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' ]