104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
"""
|
|
Database connections for Battery Service
|
|
"""
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
|
import redis.asyncio as redis
|
|
import logging
|
|
import os
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Database configuration
|
|
MONGO_URL = os.getenv("MONGO_URL", "mongodb://localhost:27017")
|
|
DATABASE_NAME = os.getenv("DATABASE_NAME", "energy_dashboard_batteries")
|
|
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():
|
|
"""Create MongoDB connection"""
|
|
global _mongo_client, _database
|
|
|
|
try:
|
|
_mongo_client = AsyncIOMotorClient(MONGO_URL)
|
|
_database = _mongo_client[DATABASE_NAME]
|
|
|
|
# Test connection
|
|
await _database.command("ping")
|
|
logger.info(f"Connected to MongoDB: {DATABASE_NAME}")
|
|
|
|
# Create indexes
|
|
await create_indexes()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to MongoDB: {e}")
|
|
raise
|
|
|
|
async def connect_to_redis():
|
|
"""Create Redis connection"""
|
|
global _redis_client
|
|
|
|
try:
|
|
_redis_client = redis.from_url(REDIS_URL, decode_responses=True)
|
|
await _redis_client.ping()
|
|
logger.info("Connected to Redis")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect to Redis: {e}")
|
|
raise
|
|
|
|
async def close_mongo_connection():
|
|
"""Close MongoDB connection"""
|
|
global _mongo_client
|
|
|
|
if _mongo_client:
|
|
_mongo_client.close()
|
|
logger.info("Disconnected from MongoDB")
|
|
|
|
async def get_database() -> AsyncIOMotorDatabase:
|
|
"""Get database instance"""
|
|
global _database
|
|
|
|
if _database is None:
|
|
raise RuntimeError("Database not initialized. Call connect_to_mongo() first.")
|
|
|
|
return _database
|
|
|
|
async def get_redis() -> redis.Redis:
|
|
"""Get Redis instance"""
|
|
global _redis_client
|
|
|
|
if _redis_client is None:
|
|
raise RuntimeError("Redis not initialized. Call connect_to_redis() first.")
|
|
|
|
return _redis_client
|
|
|
|
async def create_indexes():
|
|
"""Create database indexes for performance"""
|
|
db = await get_database()
|
|
|
|
# Indexes for batteries collection
|
|
await db.batteries.create_index("battery_id", unique=True)
|
|
await db.batteries.create_index("state")
|
|
await db.batteries.create_index("building")
|
|
await db.batteries.create_index("room")
|
|
await db.batteries.create_index("last_updated")
|
|
|
|
# Indexes for battery_history collection
|
|
await db.battery_history.create_index([("battery_id", 1), ("timestamp", -1)])
|
|
await db.battery_history.create_index("timestamp")
|
|
|
|
# Indexes for maintenance_alerts collection
|
|
await db.maintenance_alerts.create_index([("battery_id", 1), ("alert_type", 1)])
|
|
await db.maintenance_alerts.create_index("timestamp")
|
|
await db.maintenance_alerts.create_index("severity")
|
|
|
|
# Indexes for battery_events collection
|
|
await db.battery_events.create_index([("battery_id", 1), ("timestamp", -1)])
|
|
await db.battery_events.create_index("event_type")
|
|
|
|
logger.info("Database indexes created") |