changes
This commit is contained in:
HotSwapp
409
examples/websocket_pool_example.py
Normal file
409
examples/websocket_pool_example.py
Normal file
@@ -0,0 +1,409 @@
|
||||
"""
|
||||
WebSocket Connection Pool Usage Examples
|
||||
|
||||
This file demonstrates how to use the WebSocket connection pooling system
|
||||
in the Delphi Database application.
|
||||
|
||||
Examples include:
|
||||
- Basic WebSocket endpoint with pooling
|
||||
- Custom message handling
|
||||
- Topic-based broadcasting
|
||||
- Connection monitoring
|
||||
- Admin management integration
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from typing import Set, Optional, Dict, Any
|
||||
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
from app.middleware.websocket_middleware import (
|
||||
get_websocket_manager,
|
||||
websocket_endpoint,
|
||||
websocket_auth_dependency,
|
||||
WebSocketManager
|
||||
)
|
||||
from app.services.websocket_pool import (
|
||||
WebSocketMessage,
|
||||
MessageType,
|
||||
websocket_connection
|
||||
)
|
||||
from app.models.user import User
|
||||
|
||||
|
||||
# Example 1: Basic WebSocket endpoint with automatic pooling
|
||||
@websocket_endpoint(topics={"notifications"}, require_auth=True)
|
||||
async def basic_websocket_handler(
|
||||
websocket: WebSocket,
|
||||
connection_id: str,
|
||||
manager: WebSocketManager,
|
||||
message: WebSocketMessage
|
||||
):
|
||||
"""
|
||||
Basic WebSocket handler using the pooling decorator.
|
||||
Automatically handles connection management, authentication, and cleanup.
|
||||
"""
|
||||
if message.type == "user_action":
|
||||
# Handle user actions
|
||||
response = WebSocketMessage(
|
||||
type="action_response",
|
||||
data={"status": "received", "action": message.data.get("action")}
|
||||
)
|
||||
await manager.pool._send_to_connection(connection_id, response)
|
||||
|
||||
|
||||
# Example 2: Manual WebSocket management with custom logic
|
||||
async def manual_websocket_handler(websocket: WebSocket, topic: str):
|
||||
"""
|
||||
Manual WebSocket handling with direct pool management.
|
||||
Provides more control over connection lifecycle and message handling.
|
||||
"""
|
||||
manager = get_websocket_manager()
|
||||
|
||||
# Custom message handler
|
||||
async def handle_custom_message(connection_id: str, message: WebSocketMessage):
|
||||
if message.type == "chat_message":
|
||||
# Broadcast chat message to all subscribers
|
||||
await manager.broadcast_to_topic(
|
||||
topic=topic,
|
||||
message_type="chat_broadcast",
|
||||
data={
|
||||
"user": message.data.get("user", "Anonymous"),
|
||||
"message": message.data.get("message", ""),
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
},
|
||||
exclude_connection_id=connection_id
|
||||
)
|
||||
elif message.type == "typing":
|
||||
# Broadcast typing indicator
|
||||
await manager.broadcast_to_topic(
|
||||
topic=topic,
|
||||
message_type="user_typing",
|
||||
data={
|
||||
"user": message.data.get("user", "Anonymous"),
|
||||
"typing": message.data.get("typing", False)
|
||||
},
|
||||
exclude_connection_id=connection_id
|
||||
)
|
||||
|
||||
# Handle the connection
|
||||
await manager.handle_connection(
|
||||
websocket=websocket,
|
||||
topics={topic},
|
||||
require_auth=True,
|
||||
metadata={"chat_room": topic},
|
||||
message_handler=handle_custom_message
|
||||
)
|
||||
|
||||
|
||||
# Example 3: Low-level pool usage with context manager
|
||||
async def low_level_websocket_example(websocket: WebSocket, user_id: int):
|
||||
"""
|
||||
Low-level WebSocket handling using the connection context manager directly.
|
||||
Provides maximum control over the connection lifecycle.
|
||||
"""
|
||||
await websocket.accept()
|
||||
|
||||
async with websocket_connection(
|
||||
websocket=websocket,
|
||||
user_id=user_id,
|
||||
topics={"user_updates"},
|
||||
metadata={"example": "low_level"}
|
||||
) as (connection_id, pool):
|
||||
|
||||
# Send welcome message
|
||||
welcome = WebSocketMessage(
|
||||
type="welcome",
|
||||
data={
|
||||
"connection_id": connection_id,
|
||||
"message": "Connected to low-level example"
|
||||
}
|
||||
)
|
||||
await pool._send_to_connection(connection_id, welcome)
|
||||
|
||||
# Handle messages manually
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
data = await websocket.receive_text()
|
||||
|
||||
# Parse and handle message
|
||||
import json
|
||||
message_dict = json.loads(data)
|
||||
message = WebSocketMessage(**message_dict)
|
||||
|
||||
if message.type == "ping":
|
||||
pong = WebSocketMessage(type="pong", data={"timestamp": message.timestamp})
|
||||
await pool._send_to_connection(connection_id, pong)
|
||||
|
||||
elif message.type == "echo":
|
||||
echo = WebSocketMessage(
|
||||
type="echo_response",
|
||||
data={"original": message.data, "echoed_at": datetime.now(timezone.utc).isoformat()}
|
||||
)
|
||||
await pool._send_to_connection(connection_id, echo)
|
||||
|
||||
except WebSocketDisconnect:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error handling message: {e}")
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print(f"Connection error: {e}")
|
||||
|
||||
|
||||
# Example 4: Broadcasting service
|
||||
class NotificationBroadcaster:
|
||||
"""
|
||||
Service for broadcasting notifications to different user groups.
|
||||
Demonstrates how to use the pool for system-wide notifications.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.manager = get_websocket_manager()
|
||||
|
||||
async def broadcast_system_announcement(self, message: str, priority: str = "info"):
|
||||
"""Broadcast system announcement to all connected users"""
|
||||
sent_count = await self.manager.broadcast_to_topic(
|
||||
topic="system_announcements",
|
||||
message_type="system_announcement",
|
||||
data={
|
||||
"message": message,
|
||||
"priority": priority,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
)
|
||||
return sent_count
|
||||
|
||||
async def notify_user_group(self, group: str, notification_type: str, data: Dict[str, Any]):
|
||||
"""Send notification to a specific user group"""
|
||||
topic = f"group_{group}"
|
||||
sent_count = await self.manager.broadcast_to_topic(
|
||||
topic=topic,
|
||||
message_type=notification_type,
|
||||
data=data
|
||||
)
|
||||
return sent_count
|
||||
|
||||
async def send_personal_notification(self, user_id: int, notification_type: str, data: Dict[str, Any]):
|
||||
"""Send personal notification to a specific user"""
|
||||
sent_count = await self.manager.send_to_user(
|
||||
user_id=user_id,
|
||||
message_type=notification_type,
|
||||
data=data
|
||||
)
|
||||
return sent_count
|
||||
|
||||
|
||||
# Example 5: Connection monitoring and health checks
|
||||
class ConnectionMonitor:
|
||||
"""
|
||||
Service for monitoring WebSocket connections and health.
|
||||
Demonstrates how to use the pool for system monitoring.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.manager = get_websocket_manager()
|
||||
|
||||
async def get_connection_stats(self) -> Dict[str, Any]:
|
||||
"""Get comprehensive connection statistics"""
|
||||
return await self.manager.get_stats()
|
||||
|
||||
async def health_check_all_connections(self) -> Dict[str, Any]:
|
||||
"""Perform health check on all connections"""
|
||||
pool = self.manager.pool
|
||||
|
||||
async with pool._connections_lock:
|
||||
connection_ids = list(pool._connections.keys())
|
||||
|
||||
healthy = 0
|
||||
stale = 0
|
||||
total = len(connection_ids)
|
||||
|
||||
for connection_id in connection_ids:
|
||||
connection_info = await pool.get_connection_info(connection_id)
|
||||
if connection_info:
|
||||
if connection_info.is_alive():
|
||||
healthy += 1
|
||||
if connection_info.is_stale():
|
||||
stale += 1
|
||||
|
||||
return {
|
||||
"total_connections": total,
|
||||
"healthy_connections": healthy,
|
||||
"stale_connections": stale,
|
||||
"health_percentage": (healthy / total * 100) if total > 0 else 100
|
||||
}
|
||||
|
||||
async def cleanup_stale_connections(self) -> int:
|
||||
"""Manually cleanup stale connections"""
|
||||
pool = self.manager.pool
|
||||
stats_before = await pool.get_stats()
|
||||
await pool._cleanup_stale_connections()
|
||||
stats_after = await pool.get_stats()
|
||||
return stats_before["active_connections"] - stats_after["active_connections"]
|
||||
|
||||
|
||||
# Example 6: Real-time data streaming
|
||||
class RealTimeDataStreamer:
|
||||
"""
|
||||
Service for streaming real-time data to WebSocket clients.
|
||||
Demonstrates how to use the pool for continuous data updates.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.manager = get_websocket_manager()
|
||||
self._streaming_tasks: Dict[str, asyncio.Task] = {}
|
||||
|
||||
async def start_data_stream(self, topic: str, data_source: callable, interval: float = 1.0):
|
||||
"""Start streaming data to a topic"""
|
||||
if topic in self._streaming_tasks:
|
||||
return False # Already streaming
|
||||
|
||||
async def stream_data():
|
||||
while True:
|
||||
try:
|
||||
# Get data from source
|
||||
data = await data_source() if asyncio.iscoroutinefunction(data_source) else data_source()
|
||||
|
||||
# Broadcast to subscribers
|
||||
await self.manager.broadcast_to_topic(
|
||||
topic=topic,
|
||||
message_type="data_update",
|
||||
data={
|
||||
"data": data,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat()
|
||||
}
|
||||
)
|
||||
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"Error in data stream {topic}: {e}")
|
||||
await asyncio.sleep(interval * 2) # Back off on error
|
||||
|
||||
task = asyncio.create_task(stream_data())
|
||||
self._streaming_tasks[topic] = task
|
||||
return True
|
||||
|
||||
async def stop_data_stream(self, topic: str):
|
||||
"""Stop streaming data to a topic"""
|
||||
if topic in self._streaming_tasks:
|
||||
self._streaming_tasks[topic].cancel()
|
||||
del self._streaming_tasks[topic]
|
||||
return True
|
||||
return False
|
||||
|
||||
async def stop_all_streams(self):
|
||||
"""Stop all data streams"""
|
||||
for task in self._streaming_tasks.values():
|
||||
task.cancel()
|
||||
self._streaming_tasks.clear()
|
||||
|
||||
|
||||
# Example FastAPI application demonstrating usage
|
||||
def create_example_app() -> FastAPI:
|
||||
"""Create example FastAPI application with WebSocket endpoints"""
|
||||
app = FastAPI(title="WebSocket Pool Example")
|
||||
|
||||
# Initialize services
|
||||
broadcaster = NotificationBroadcaster()
|
||||
monitor = ConnectionMonitor()
|
||||
streamer = RealTimeDataStreamer()
|
||||
|
||||
@app.websocket("/ws/basic")
|
||||
async def basic_endpoint(websocket: WebSocket):
|
||||
"""Basic WebSocket endpoint with automatic pooling"""
|
||||
await basic_websocket_handler(websocket, "basic", get_websocket_manager(), None)
|
||||
|
||||
@app.websocket("/ws/chat/{room}")
|
||||
async def chat_endpoint(websocket: WebSocket, room: str):
|
||||
"""Chat room WebSocket endpoint"""
|
||||
await manual_websocket_handler(websocket, f"chat_{room}")
|
||||
|
||||
@app.websocket("/ws/user/{user_id}")
|
||||
async def user_endpoint(websocket: WebSocket, user_id: int):
|
||||
"""User-specific WebSocket endpoint"""
|
||||
await low_level_websocket_example(websocket, user_id)
|
||||
|
||||
@app.get("/")
|
||||
async def index():
|
||||
"""Simple HTML page for testing WebSocket connections"""
|
||||
return HTMLResponse("""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>WebSocket Pool Example</title>
|
||||
</head>
|
||||
<body>
|
||||
<h1>WebSocket Pool Example</h1>
|
||||
<div id="messages"></div>
|
||||
<input type="text" id="messageInput" placeholder="Type a message...">
|
||||
<button onclick="sendMessage()">Send</button>
|
||||
|
||||
<script>
|
||||
const ws = new WebSocket('ws://localhost:8000/ws/basic');
|
||||
const messages = document.getElementById('messages');
|
||||
|
||||
ws.onmessage = function(event) {
|
||||
const message = JSON.parse(event.data);
|
||||
const div = document.createElement('div');
|
||||
div.textContent = JSON.stringify(message, null, 2);
|
||||
messages.appendChild(div);
|
||||
};
|
||||
|
||||
function sendMessage() {
|
||||
const input = document.getElementById('messageInput');
|
||||
const message = {
|
||||
type: 'user_action',
|
||||
data: { action: input.value }
|
||||
};
|
||||
ws.send(JSON.stringify(message));
|
||||
input.value = '';
|
||||
}
|
||||
|
||||
// Send ping every 30 seconds
|
||||
setInterval(() => {
|
||||
ws.send(JSON.stringify({ type: 'ping' }));
|
||||
}, 30000);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
""")
|
||||
|
||||
@app.post("/api/broadcast/system")
|
||||
async def broadcast_system(message: str, priority: str = "info"):
|
||||
"""Broadcast system message to all users"""
|
||||
sent_count = await broadcaster.broadcast_system_announcement(message, priority)
|
||||
return {"sent_count": sent_count}
|
||||
|
||||
@app.get("/api/monitor/stats")
|
||||
async def get_monitor_stats():
|
||||
"""Get connection monitoring statistics"""
|
||||
return await monitor.get_connection_stats()
|
||||
|
||||
@app.get("/api/monitor/health")
|
||||
async def get_health_status():
|
||||
"""Get connection health status"""
|
||||
return await monitor.health_check_all_connections()
|
||||
|
||||
@app.post("/api/monitor/cleanup")
|
||||
async def cleanup_connections():
|
||||
"""Manually cleanup stale connections"""
|
||||
cleaned = await monitor.cleanup_stale_connections()
|
||||
return {"cleaned_connections": cleaned}
|
||||
|
||||
return app
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the example application
|
||||
import uvicorn
|
||||
|
||||
app = create_example_app()
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
Reference in New Issue
Block a user