Update app/services/websocket.py

app/services/websocket.py

@@ -1,63 +1,80 @@
-from typing import List, Dict, Any
-from fastapi import WebSocket
-from datetime import datetime
-from ..db.database import db
-from ..utils.cache import cache
-
-# Store active WebSocket connections
-active_connections: List[WebSocket] = []
-
-async def connect(websocket: WebSocket):
-    """Accept a new WebSocket connection"""
-    await websocket.accept()
-    active_connections.append(websocket)
-
-async def disconnect(websocket: WebSocket):
-    """Remove a WebSocket connection"""
-    if websocket in active_connections:
-        active_connections.remove(websocket)
-
-async def broadcast_message(message: dict):
-    """Broadcast a message to all connected clients"""
-    disconnected = []
-    for connection in active_connections:
-        try:
-            await connection.send_json(message)
-        except:
-            disconnected.append(connection)
-    
-    # Clean up disconnected clients
-    for connection in disconnected:
-        await disconnect(connection)
-
-async def create_and_broadcast_notification(
-    user_id: str,
-    title: str,
-    message: str,
-    notification_type: str,
-    data: Dict[str, Any] = None
-) -> Dict[str, Any]:
-    """Create and broadcast a notification"""
-    notification = {
-        "user_id": user_id,
-        "title": title,
-        "message": message,
-        "type": notification_type,
-        "data": data,
-        "created_at": datetime.utcnow(),
-        "read": False
-    }
-    
-    # Store in database
-    await db.db["notifications"].insert_one(notification)
-    
-    # Broadcast to connected clients
-    await broadcast_message({
-        "type": "notification",
-        "data": notification
-    })
-    
-    # Clear user's notification cache
-    await cache.delete_cache(f"user_notifications:{user_id}")
-    
-    return notification
+from typing import List, Dict, Any
+from fastapi import WebSocket
+from datetime import datetime
+from ..db.database import db
+from ..utils.cache import cache
+from ..db.models import Notification
+
+# Store active WebSocket connections
+active_connections: List[WebSocket] = []
+
+async def connect(websocket: WebSocket):
+    """Accept a new WebSocket connection"""
+    await websocket.accept()
+    active_connections.append(websocket)
+
+async def disconnect(websocket: WebSocket):
+    """Remove a WebSocket connection"""
+    if websocket in active_connections:
+        active_connections.remove(websocket)
+
+async def broadcast_message(message: dict):
+    """Broadcast a message to all connected clients"""
+    disconnected = []
+    for connection in active_connections:
+        try:
+            await connection.send_json(message)
+        except:
+            disconnected.append(connection)
+    
+    # Clean up disconnected clients
+    for connection in disconnected:
+        await disconnect(connection)
+
+async def create_and_broadcast_notification(
+    user_id: str,
+    title: str,
+    message: str,
+    notification_type: str,
+    data: Dict[str, Any] = None
+) -> Dict[str, Any]:
+    """Create and broadcast a notification"""
+    async with db.session() as session:
+        # Create notification
+        notification = Notification(
+            user_id=user_id,
+            title=title,
+            message=message,
+            type=notification_type,
+            data=data,
+            created_at=datetime.utcnow(),
+            read=False
+        )
+        
+        # Store in database
+        session.add(notification)
+        await session.commit()
+        await session.refresh(notification)
+        
+        # Convert to dict for broadcasting
+        notification_dict = {
+            "id": notification.id,
+            "user_id": notification.user_id,
+            "title": notification.title,
+            "message": notification.message,
+            "type": notification.type,
+            "data": notification.data,
+            "created_at": notification.created_at.isoformat(),
+            "read": notification.read
+        }
+        
+        # Broadcast to connected clients
+        await broadcast_message({
+            "type": "notification",
+            "data": notification_dict
+        })
+        
+        # Clear user's notification cache
+        await cache.delete_cache(f"user_notifications:{user_id}")
+        
+        return notification_dict