bakery-ia/services/data/app/core/database.py

"""
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'
]