Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from typing import Dict | |
| import sqlite3 | |
| import json | |
| import os | |
| app = FastAPI() | |
| # Enable CORS so your mobile app/web app can connect without security blocks | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| # Your mounted storage bucket path | |
| DB_PATH = "/data/lowkey.db" | |
| def init_db(): | |
| # Ensures the database file is created directly inside your 1GB persistent bucket | |
| conn = sqlite3.connect(DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute(''' | |
| CREATE TABLE IF NOT EXISTS queue ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| sender TEXT, | |
| receiver TEXT, | |
| message TEXT, | |
| timestamp DATETIME DEFAULT CURRENT_TIMESTAMP | |
| ) | |
| ''') | |
| conn.commit() | |
| conn.close() | |
| # Initialize DB on server startup | |
| async def startup_event(): | |
| init_db() | |
| class ConnectionManager: | |
| def __init__(self): | |
| self.active_connections: Dict[str, WebSocket] = {} | |
| async def connect(self, user_id: str, websocket: WebSocket): | |
| await websocket.accept() | |
| self.active_connections[user_id] = websocket | |
| await self.flush_offline_messages(user_id) | |
| def disconnect(self, user_id: str): | |
| if user_id in self.active_connections: | |
| del self.active_connections[user_id] | |
| async def send_message(self, receiver_id: str, message: dict): | |
| if receiver_id in self.active_connections: | |
| try: | |
| await self.active_connections[receiver_id].send_text(json.dumps(message)) | |
| return True | |
| except Exception: | |
| # If sending fails halfway through, fallback to saving offline | |
| self.queue_message(message['sender'], receiver_id, message['text']) | |
| return False | |
| else: | |
| self.queue_message(message['sender'], receiver_id, message['text']) | |
| return False | |
| def queue_message(self, sender, receiver, text): | |
| conn = sqlite3.connect(DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("INSERT INTO queue (sender, receiver, message) VALUES (?, ?, ?)", (sender, receiver, text)) | |
| conn.commit() | |
| conn.close() | |
| async def flush_offline_messages(self, user_id: str): | |
| conn = sqlite3.connect(DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT id, sender, message FROM queue WHERE receiver = ?", (user_id,)) | |
| rows = cursor.fetchall() | |
| if rows: | |
| websocket = self.active_connections[user_id] | |
| for row in rows: | |
| msg_id, sender, text = row | |
| payload = {"sender": sender, "text": text, "status": "offline_delivery"} | |
| try: | |
| await websocket.send_text(json.dumps(payload)) | |
| except Exception: | |
| conn.close() | |
| return # Stop if user disconnects mid-delivery | |
| cursor.execute("DELETE FROM queue WHERE receiver = ?", (user_id,)) | |
| conn.commit() | |
| conn.close() | |
| manager = ConnectionManager() | |
| def read_root(): | |
| return {"status": "Lowkey Server is up and running 🔥"} | |
| async def websocket_endpoint(websocket: WebSocket, user_id: str): | |
| await manager.connect(user_id, websocket) | |
| try: | |
| while True: | |
| data = await websocket.receive_text() | |
| payload = json.loads(data) | |
| # Action: App client says "I downloaded and saved this safely" | |
| if payload.get("action") == "seen_and_saved": | |
| continue | |
| target_receiver = payload.get("receiver") | |
| msg_data = { | |
| "sender": user_id, | |
| "text": payload.get("text") | |
| } | |
| await manager.send_message(target_receiver, msg_data) | |
| except WebSocketDisconnect: | |
| manager.disconnect(user_id) |