first commit

main.py

@@ -0,0 +1,202 @@
+
+import asyncio
+import json
+import redis.asyncio as redis
+import time
+import os
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Depends, Query
+from fastapi.middleware.cors import CORSMiddleware
+from typing import List, Optional
+import logging
+from contextlib import asynccontextmanager
+
+# Import our custom modules
+from database import connect_to_mongo, close_mongo_connection, redis_manager, schedule_cleanup
+from persistence import persistence_service
+from models import DataQuery, DataResponse, HealthCheck
+from api import router as api_router
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Application startup time for uptime calculation
+app_start_time = time.time()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager"""
+    # Startup
+    logger.info("Application starting up...")
+    
+    # Connect to databases
+    await connect_to_mongo()
+    await persistence_service.initialize()
+    
+    # Start background tasks
+    asyncio.create_task(redis_subscriber())
+    asyncio.create_task(schedule_cleanup())
+    
+    logger.info("Application startup complete")
+    
+    yield
+    
+    # Shutdown
+    logger.info("Application shutting down...")
+    await close_mongo_connection()
+    await redis_manager.disconnect()
+    logger.info("Application shutdown complete")
+
+app = FastAPI(
+    title="Energy Monitoring Dashboard API",
+    description="Real-time energy monitoring and IoT sensor data management system",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Configure appropriately for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Include API router
+app.include_router(api_router, prefix="/api/v1")
+
+# In-memory store for active WebSocket connections
+active_connections: List[WebSocket] = []
+
+# Redis channel to subscribe to
+REDIS_CHANNEL = "energy_data"
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """
+    WebSocket endpoint that connects a client, adds them to the active pool,
+    and removes them on disconnection.
+    """
+    await websocket.accept()
+    active_connections.append(websocket)
+    logger.info(f"New client connected. Total clients: {len(active_connections)}")
+    try:
+        while True:
+            # Keep the connection alive
+            await websocket.receive_text()
+    except WebSocketDisconnect:
+        active_connections.remove(websocket)
+        logger.info(f"Client disconnected. Total clients: {len(active_connections)}")
+
+
+async def redis_subscriber():
+    """
+    Connects to Redis, subscribes to the specified channel, and broadcasts
+    messages to all active WebSocket clients. Also persists data to MongoDB.
+    """
+    logger.info("Starting Redis subscriber...")
+    try:
+        r = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
+        await r.ping()
+        logger.info("Successfully connected to Redis for subscription.")
+    except Exception as e:
+        logger.error(f"Could not connect to Redis for subscription: {e}")
+        return
+
+    pubsub = r.pubsub()
+    await pubsub.subscribe(REDIS_CHANNEL)
+    
+    logger.info(f"Subscribed to Redis channel: '{REDIS_CHANNEL}'")
+    while True:
+        try:
+            message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
+            if message:
+                message_data = message['data']
+                logger.debug(f"Received from Redis: {message_data}")
+                
+                # Process and persist the data
+                await persistence_service.process_sensor_message(message_data)
+                
+                # Broadcast message to all connected WebSocket clients
+                if active_connections:
+                    await asyncio.gather(
+                        *[connection.send_text(message_data) for connection in active_connections],
+                        return_exceptions=True
+                    )
+                
+        except Exception as e:
+            logger.error(f"Error in Redis subscriber loop: {e}")
+            # Add a delay to prevent rapid-fire errors
+            await asyncio.sleep(5)
+
+
+@app.get("/")
+async def read_root():
+    """Root endpoint with basic system information"""
+    return {
+        "message": "Energy Monitoring Dashboard Backend",
+        "version": "1.0.0",
+        "status": "running",
+        "uptime_seconds": time.time() - app_start_time
+    }
+
+
+@app.get("/health", response_model=HealthCheck)
+async def health_check():
+    """Health check endpoint"""
+    try:
+        # Check database connections
+        mongodb_connected = True
+        redis_connected = True
+        
+        try:
+            await persistence_service.db.command("ping")
+        except:
+            mongodb_connected = False
+            
+        try:
+            await redis_manager.redis_client.ping()
+        except:
+            redis_connected = False
+        
+        # Get system statistics
+        stats = await persistence_service.get_sensor_statistics()
+        
+        # Determine overall status
+        status = "healthy"
+        if not mongodb_connected or not redis_connected:
+            status = "degraded"
+        
+        return HealthCheck(
+            status=status,
+            mongodb_connected=mongodb_connected,
+            redis_connected=redis_connected,
+            total_sensors=stats.get("total_sensors", 0),
+            active_sensors=stats.get("active_sensors", 0),
+            total_readings=stats.get("total_readings", 0),
+            uptime_seconds=time.time() - app_start_time
+        )
+        
+    except Exception as e:
+        logger.error(f"Health check failed: {e}")
+        raise HTTPException(status_code=503, detail="Service Unavailable")
+
+
+@app.get("/status")
+async def system_status():
+    """Detailed system status endpoint"""
+    try:
+        stats = await persistence_service.get_sensor_statistics()
+        
+        return {
+            "timestamp": time.time(),
+            "uptime_seconds": time.time() - app_start_time,
+            "active_websocket_connections": len(active_connections),
+            "database_stats": stats
+        }
+        
+    except Exception as e:
+        logger.error(f"Status check failed: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")