Redis + Websocket

app.py

@@ -8,11 +8,12 @@ from bson.objectid import ObjectId
 import redis
 import json
 from redis.exceptions import RedisError
-from pydantic import BaseModel
+from pydantic import BaseModel, validator
+from pymongo.errors import ConnectionFailure
 
 # Initialize FastAPI app
 app = FastAPI()
-socket_manager = SocketManager(app=app)
+socket_manager = SocketManager(app=app, cors_allowed_origins="*", mount_location="/socket.io")
 
 # Add CORS middleware
 app.add_middleware(
@@ -34,6 +35,14 @@ db = client["test"]
 pdf_collection = db["PdfDetails"]
 voter_collection = db["Voters"]
 
+# Kiểm tra kết nối MongoDB khi khởi động
+try:
+    # The ismaster command is cheap and does not require auth.
+    client.admin.command('ismaster')
+except ConnectionFailure:
+    print("Server not available")
+    raise
+
 # Pydantic models for request validation
 class VoterRegistration(BaseModel):
     name: str
@@ -45,6 +54,13 @@ class VoteRequest(BaseModel):
     file_id: str
     vote_count: Optional[int] = 1
 
+    class Config:
+        @validator('vote_count')
+        def validate_vote_count(cls, v):
+            if v <= 0:
+                raise ValueError('Vote count must be greater than 0')
+            return v
+
 @app.post("/upload-files")
 async def upload_file(
     title: str = Form(...),
@@ -141,35 +157,60 @@ async def vote_by_voter(vote_request: VoteRequest):
 
 @app.get("/get-voter")
 async def get_voter(voter_id: str):
-    # Check cache
-    cached_voter = redis_client.get(f'voter:{voter_id}')
-    if cached_voter:
-        return json.loads(cached_voter)
-
-    voter = voter_collection.find_one({'_id': ObjectId(voter_id)})
-    if not voter:
-        raise HTTPException(status_code=404, detail="Voter not found")
-    
-    voter_data = {
-        'status': 'ok',
-        'name': voter['name'],
-        'group': voter['group'],
-        'role': voter['role'],
-        'number_of_votes': voter['number_of_votes']
-    }
-    
-    # Save to cache
-    redis_client.setex(f'voter:{voter_id}', CACHE_EXPIRE_TIME, json.dumps(voter_data))
-    return voter_data
+    try:
+        # Check cache
+        try:
+            cached_voter = redis_client.get(f'voter:{voter_id}')
+            if cached_voter:
+                return json.loads(cached_voter)
+        except RedisError as redis_err:
+            print(f"Redis error: {redis_err}")
+            # Continue without cache if Redis fails
+            pass
+
+        voter = voter_collection.find_one({'_id': ObjectId(voter_id)})
+        if not voter:
+            raise HTTPException(status_code=404, detail="Voter not found")
+        
+        voter_data = {
+            'status': 'ok',
+            'name': voter['name'],
+            'group': voter['group'],
+            'role': voter['role'],
+            'number_of_votes': voter['number_of_votes']
+        }
+        
+        # Try to save to cache
+        try:
+            redis_client.setex(f'voter:{voter_id}', CACHE_EXPIRE_TIME, json.dumps(voter_data))
+        except RedisError as redis_err:
+            print(f"Redis caching error: {redis_err}")
+            # Continue even if caching fails
+            pass
+
+        return voter_data
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 @socket_manager.on('connect')
-async def handle_connect():
-    print('Client connected')
+async def handle_connect(sid, environ):
+    print(f'Client connected: {sid}')
 
 @socket_manager.on('disconnect')
-async def handle_disconnect():
-    print('Client disconnected')
+async def handle_disconnect(sid):
+    print(f'Client disconnected: {sid}')
+
+@app.get("/")
+async def index():
+    return {'status': 'ok', 'message': 'Server is running'}
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    # Close MongoDB connection
+    client.close()
+    # Close Redis connection
+    redis_client.close()
 
-if __name__ == "__main__":
+if _name_ == "_main_":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=5000)