sac4cps-backend/microservices/demand-response-service/database.py

"""
Database configuration and connection management for Demand Response Service
"""

from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
import redis.asyncio as redis
import logging
import os

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configuration from environment variables
MONGO_URL = os.getenv("MONGO_URL", "mongodb://localhost:27017")
DATABASE_NAME = os.getenv("DATABASE_NAME", "energy_dashboard_demand_response")
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")

# Global database clients
_mongo_client: AsyncIOMotorClient = None
_database: AsyncIOMotorDatabase = None
_redis_client: redis.Redis = None


async def connect_to_mongo():
    """Initialize MongoDB connection and create indexes"""
    global _mongo_client, _database

    try:
        logger.info(f"Connecting to MongoDB at {MONGO_URL}")
        _mongo_client = AsyncIOMotorClient(MONGO_URL)
        _database = _mongo_client[DATABASE_NAME]

        # Test connection
        await _database.command("ping")
        logger.info(f"Successfully connected to MongoDB database: {DATABASE_NAME}")

        # Create indexes
        await create_indexes()

    except Exception as e:
        logger.error(f"Failed to connect to MongoDB: {e}")
        raise


async def close_mongo_connection():
    """Close MongoDB connection"""
    global _mongo_client

    if _mongo_client:
        _mongo_client.close()
        logger.info("MongoDB connection closed")


async def get_database() -> AsyncIOMotorDatabase:
    """Get database instance"""
    if _database is None:
        await connect_to_mongo()
    return _database


async def connect_to_redis():
    """Initialize Redis connection"""
    global _redis_client

    try:
        logger.info(f"Connecting to Redis at {REDIS_URL}")
        _redis_client = redis.from_url(REDIS_URL, decode_responses=True)

        # Test connection
        await _redis_client.ping()
        logger.info("Successfully connected to Redis")

    except Exception as e:
        logger.error(f"Failed to connect to Redis: {e}")
        raise


async def close_redis_connection():
    """Close Redis connection"""
    global _redis_client

    if _redis_client:
        await _redis_client.close()
        logger.info("Redis connection closed")


async def get_redis() -> redis.Redis:
    """Get Redis client instance"""
    if _redis_client is None:
        await connect_to_redis()
    return _redis_client


async def create_indexes():
    """Create MongoDB indexes for optimal query performance"""
    db = await get_database()

    logger.info("Creating MongoDB indexes...")

    try:
        # Indexes for demand_response_invitations collection
        await db.demand_response_invitations.create_index("event_id", unique=True)
        await db.demand_response_invitations.create_index([("event_time", 1), ("status", 1)])
        await db.demand_response_invitations.create_index("status")
        await db.demand_response_invitations.create_index("created_at")
        await db.demand_response_invitations.create_index("response")
        logger.info("Created indexes for demand_response_invitations collection")

        # Indexes for demand_response_events collection
        await db.demand_response_events.create_index("event_id", unique=True)
        await db.demand_response_events.create_index([("start_time", 1), ("status", 1)])
        await db.demand_response_events.create_index([("status", 1), ("start_time", 1)])
        await db.demand_response_events.create_index("status")
        await db.demand_response_events.create_index("invitation_id")
        logger.info("Created indexes for demand_response_events collection")

        # Indexes for demand_response_responses collection
        await db.demand_response_responses.create_index([("event_id", 1), ("device_id", 1)], unique=True)
        await db.demand_response_responses.create_index("event_id")
        await db.demand_response_responses.create_index("device_id")
        await db.demand_response_responses.create_index("responded_at")
        logger.info("Created indexes for demand_response_responses collection")

        # Indexes for flexibility_snapshots collection (with TTL for auto-cleanup)
        await db.flexibility_snapshots.create_index([("timestamp", -1)])
        await db.flexibility_snapshots.create_index(
            "timestamp",
            expireAfterSeconds=7776000  # 90 days TTL
        )
        logger.info("Created indexes for flexibility_snapshots collection")

        # Indexes for auto_response_config collection (singleton document)
        await db.auto_response_config.create_index("config_id", unique=True)
        logger.info("Created indexes for auto_response_config collection")

        # Indexes for device_instructions collection
        await db.device_instructions.create_index("device_id", unique=True)
        await db.device_instructions.create_index("updated_at")
        logger.info("Created indexes for device_instructions collection")

        logger.info("All MongoDB indexes created successfully")

    except Exception as e:
        logger.error(f"Error creating indexes: {e}")
        # Don't raise - indexes may already exist


async def initialize_default_config():
    """Initialize default auto-response configuration if it doesn't exist"""
    db = await get_database()

    try:
        # Check if default config exists
        existing_config = await db.auto_response_config.find_one({"config_id": "default"})

        if not existing_config:
            default_config = {
                "config_id": "default",
                "enabled": False,
                "max_reduction_percentage": 20.0,
                "response_delay_seconds": 300,
                "min_notice_minutes": 60,
                "created_at": None,
                "updated_at": None
            }

            await db.auto_response_config.insert_one(default_config)
            logger.info("Created default auto-response configuration")
        else:
            logger.info("Auto-response configuration already exists")

    except Exception as e:
        logger.error(f"Error initializing default config: {e}")


# Utility functions for common database operations

async def get_collection(collection_name: str):
    """Get a collection by name"""
    db = await get_database()
    return db[collection_name]


async def health_check() -> dict:
    """Check database connections health"""
    status = {
        "mongodb": False,
        "redis": False
    }

    try:
        # Check MongoDB
        db = await get_database()
        await db.command("ping")
        status["mongodb"] = True
    except Exception as e:
        logger.error(f"MongoDB health check failed: {e}")

    try:
        # Check Redis
        redis_client = await get_redis()
        await redis_client.ping()
        status["redis"] = True
    except Exception as e:
        logger.error(f"Redis health check failed: {e}")

    return status