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