first commit
This commit is contained in:
rafaeldpsilva
220
database.py
Normal file
220
database.py
Normal file
@@ -0,0 +1,220 @@
|
||||
import os
|
||||
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
|
||||
from pymongo import IndexModel, ASCENDING, DESCENDING
|
||||
from typing import Optional
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
import logging
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class MongoDB:
|
||||
client: Optional[AsyncIOMotorClient] = None
|
||||
database: Optional[AsyncIOMotorDatabase] = None
|
||||
|
||||
# Global MongoDB instance
|
||||
mongodb = MongoDB()
|
||||
|
||||
async def connect_to_mongo():
|
||||
"""Create database connection"""
|
||||
try:
|
||||
# MongoDB connection string - default to localhost for development
|
||||
mongodb_url = os.getenv("MONGODB_URL", "mongodb://localhost:27017")
|
||||
database_name = os.getenv("DATABASE_NAME", "energy_monitoring")
|
||||
|
||||
logger.info(f"Connecting to MongoDB at: {mongodb_url}")
|
||||
|
||||
# Create async MongoDB client
|
||||
mongodb.client = AsyncIOMotorClient(mongodb_url)
|
||||
|
||||
# Test the connection
|
||||
await mongodb.client.admin.command('ping')
|
||||
logger.info("Successfully connected to MongoDB")
|
||||
|
||||
# Get database
|
||||
mongodb.database = mongodb.client[database_name]
|
||||
|
||||
# Create indexes for better performance
|
||||
await create_indexes()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error connecting to MongoDB: {e}")
|
||||
raise
|
||||
|
||||
async def close_mongo_connection():
|
||||
"""Close database connection"""
|
||||
if mongodb.client:
|
||||
mongodb.client.close()
|
||||
logger.info("Disconnected from MongoDB")
|
||||
|
||||
async def create_indexes():
|
||||
"""Create database indexes for optimal performance"""
|
||||
try:
|
||||
# Sensor readings collection indexes
|
||||
sensor_readings_indexes = [
|
||||
IndexModel([("sensor_id", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
IndexModel([("timestamp", DESCENDING)]),
|
||||
IndexModel([("room", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
IndexModel([("sensor_type", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
IndexModel([("created_at", DESCENDING)]),
|
||||
]
|
||||
await mongodb.database.sensor_readings.create_indexes(sensor_readings_indexes)
|
||||
|
||||
# Room metrics collection indexes
|
||||
room_metrics_indexes = [
|
||||
IndexModel([("room", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
IndexModel([("timestamp", DESCENDING)]),
|
||||
IndexModel([("created_at", DESCENDING)]),
|
||||
]
|
||||
await mongodb.database.room_metrics.create_indexes(room_metrics_indexes)
|
||||
|
||||
# Sensor metadata collection indexes
|
||||
sensor_metadata_indexes = [
|
||||
IndexModel([("sensor_id", ASCENDING)], unique=True),
|
||||
IndexModel([("room", ASCENDING)]),
|
||||
IndexModel([("sensor_type", ASCENDING)]),
|
||||
IndexModel([("status", ASCENDING)]),
|
||||
]
|
||||
await mongodb.database.sensor_metadata.create_indexes(sensor_metadata_indexes)
|
||||
|
||||
# System events collection indexes
|
||||
system_events_indexes = [
|
||||
IndexModel([("timestamp", DESCENDING)]),
|
||||
IndexModel([("event_type", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
IndexModel([("severity", ASCENDING), ("timestamp", DESCENDING)]),
|
||||
]
|
||||
await mongodb.database.system_events.create_indexes(system_events_indexes)
|
||||
|
||||
logger.info("Database indexes created successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating indexes: {e}")
|
||||
|
||||
async def get_database() -> AsyncIOMotorDatabase:
|
||||
"""Get database instance"""
|
||||
if not mongodb.database:
|
||||
await connect_to_mongo()
|
||||
return mongodb.database
|
||||
|
||||
class RedisManager:
|
||||
"""Redis connection and operations manager"""
|
||||
|
||||
def __init__(self):
|
||||
self.redis_client = None
|
||||
self.redis_host = os.getenv("REDIS_HOST", "localhost")
|
||||
self.redis_port = int(os.getenv("REDIS_PORT", "6379"))
|
||||
self.redis_db = int(os.getenv("REDIS_DB", "0"))
|
||||
|
||||
async def connect(self):
|
||||
"""Connect to Redis"""
|
||||
try:
|
||||
import redis.asyncio as redis
|
||||
self.redis_client = redis.Redis(
|
||||
host=self.redis_host,
|
||||
port=self.redis_port,
|
||||
db=self.redis_db,
|
||||
decode_responses=True
|
||||
)
|
||||
await self.redis_client.ping()
|
||||
logger.info("Successfully connected to Redis")
|
||||
except Exception as e:
|
||||
logger.error(f"Error connecting to Redis: {e}")
|
||||
raise
|
||||
|
||||
async def disconnect(self):
|
||||
"""Disconnect from Redis"""
|
||||
if self.redis_client:
|
||||
await self.redis_client.close()
|
||||
logger.info("Disconnected from Redis")
|
||||
|
||||
async def set_sensor_data(self, sensor_id: str, data: dict, expire_time: int = 3600):
|
||||
"""Store latest sensor data in Redis with expiration"""
|
||||
if not self.redis_client:
|
||||
await self.connect()
|
||||
|
||||
key = f"sensor:latest:{sensor_id}"
|
||||
await self.redis_client.setex(key, expire_time, str(data))
|
||||
|
||||
async def get_sensor_data(self, sensor_id: str) -> Optional[dict]:
|
||||
"""Get latest sensor data from Redis"""
|
||||
if not self.redis_client:
|
||||
await self.connect()
|
||||
|
||||
key = f"sensor:latest:{sensor_id}"
|
||||
data = await self.redis_client.get(key)
|
||||
if data:
|
||||
import json
|
||||
return json.loads(data)
|
||||
return None
|
||||
|
||||
async def set_room_metrics(self, room: str, metrics: dict, expire_time: int = 1800):
|
||||
"""Store room aggregated metrics in Redis"""
|
||||
if not self.redis_client:
|
||||
await self.connect()
|
||||
|
||||
key = f"room:metrics:{room}"
|
||||
await self.redis_client.setex(key, expire_time, str(metrics))
|
||||
|
||||
async def get_room_metrics(self, room: str) -> Optional[dict]:
|
||||
"""Get room aggregated metrics from Redis"""
|
||||
if not self.redis_client:
|
||||
await self.connect()
|
||||
|
||||
key = f"room:metrics:{room}"
|
||||
data = await self.redis_client.get(key)
|
||||
if data:
|
||||
import json
|
||||
return json.loads(data)
|
||||
return None
|
||||
|
||||
async def get_all_active_sensors(self) -> list:
|
||||
"""Get list of all sensors with recent data in Redis"""
|
||||
if not self.redis_client:
|
||||
await self.connect()
|
||||
|
||||
keys = await self.redis_client.keys("sensor:latest:*")
|
||||
return [key.replace("sensor:latest:", "") for key in keys]
|
||||
|
||||
# Global Redis manager instance
|
||||
redis_manager = RedisManager()
|
||||
|
||||
async def cleanup_old_data():
|
||||
"""Cleanup old data from MongoDB (retention policy)"""
|
||||
try:
|
||||
db = await get_database()
|
||||
|
||||
# Delete sensor readings older than 90 days
|
||||
retention_date = datetime.utcnow() - timedelta(days=90)
|
||||
result = await db.sensor_readings.delete_many({
|
||||
"created_at": {"$lt": retention_date}
|
||||
})
|
||||
|
||||
if result.deleted_count > 0:
|
||||
logger.info(f"Deleted {result.deleted_count} old sensor readings")
|
||||
|
||||
# Delete room metrics older than 30 days
|
||||
retention_date = datetime.utcnow() - timedelta(days=30)
|
||||
result = await db.room_metrics.delete_many({
|
||||
"created_at": {"$lt": retention_date}
|
||||
})
|
||||
|
||||
if result.deleted_count > 0:
|
||||
logger.info(f"Deleted {result.deleted_count} old room metrics")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up old data: {e}")
|
||||
|
||||
# Scheduled cleanup task
|
||||
async def schedule_cleanup():
|
||||
"""Schedule periodic cleanup of old data"""
|
||||
while True:
|
||||
try:
|
||||
await cleanup_old_data()
|
||||
# Wait 24 hours before next cleanup
|
||||
await asyncio.sleep(24 * 60 * 60)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in scheduled cleanup: {e}")
|
||||
# Wait 1 hour before retrying
|
||||
await asyncio.sleep(60 * 60)
|
||||
Reference in New Issue
Block a user