forked from erp-dev/erp
132 lines
4.4 KiB
Python
132 lines
4.4 KiB
Python
from django.http import StreamingHttpResponse, HttpResponse
|
||
from django.views.decorators.csrf import csrf_exempt
|
||
from django.views.decorators.http import require_http_methods
|
||
from rest_framework.decorators import api_view, permission_classes
|
||
from rest_framework.permissions import AllowAny
|
||
from rest_framework.response import Response
|
||
from rest_framework import status
|
||
from . import services
|
||
import queue
|
||
import json
|
||
|
||
|
||
@csrf_exempt
|
||
@require_http_methods(["GET", "OPTIONS"])
|
||
def create_sse_event(request):
|
||
"""
|
||
创建一个 SSE 事件流响应。
|
||
客户端连接到此端点后会保持长连接，等待服务器推送事件。
|
||
|
||
使用方法：
|
||
- GET /sse/
|
||
- 保持连接打开以接收实时事件
|
||
|
||
注意：
|
||
- 使用纯 Django 视图，不使用 DRF，避免内容协商导致的 406 错误
|
||
- SSE 需要特殊的 CORS 配置
|
||
"""
|
||
# 处理 OPTIONS 预检请求
|
||
if request.method == 'OPTIONS':
|
||
response = HttpResponse()
|
||
origin = request.META.get('HTTP_ORIGIN')
|
||
if origin:
|
||
response['Access-Control-Allow-Origin'] = origin
|
||
response['Access-Control-Allow-Methods'] = 'GET, OPTIONS'
|
||
response['Access-Control-Allow-Headers'] = 'authorization, content-type, cache-control, accept'
|
||
response['Access-Control-Allow-Credentials'] = 'true'
|
||
response['Access-Control-Max-Age'] = '86400' # 24小时
|
||
return response
|
||
|
||
def event_stream():
|
||
# 创建一个同步队列用于接收消息
|
||
conn_queue = queue.Queue(maxsize=100)
|
||
services.push_connection(conn_queue)
|
||
|
||
try:
|
||
# 发送初始连接成功消息
|
||
yield f"data: {json.dumps({'type': 'connected', 'message': 'SSE connection established'})}\n\n"
|
||
|
||
# 持续从队列中获取消息并发送给客户端
|
||
while True:
|
||
try:
|
||
# 使用同步方式等待新消息，带超时（30秒）以便发送心跳
|
||
message = conn_queue.get(timeout=30.0)
|
||
yield f"data: {json.dumps(message)}\n\n"
|
||
except queue.Empty:
|
||
# 30秒超时，发送心跳保持连接活跃
|
||
yield f": heartbeat\n\n"
|
||
except Exception as e:
|
||
print(f"Error in SSE stream: {e}")
|
||
break
|
||
finally:
|
||
# 清理：从连接集合中移除此队列
|
||
services.remove_connection(conn_queue)
|
||
|
||
response = StreamingHttpResponse(
|
||
event_stream(),
|
||
content_type='text/event-stream',
|
||
)
|
||
|
||
# SSE 必需的响应头
|
||
response['Cache-Control'] = 'no-cache'
|
||
response['X-Accel-Buffering'] = 'no' # 禁用 nginx 缓冲
|
||
|
||
# CORS 响应头（django-cors-headers 中间件会自动添加，但我们显式设置以确保）
|
||
# 如果请求带有 Origin 头，手动添加 CORS 响应头
|
||
origin = request.META.get('HTTP_ORIGIN')
|
||
if origin:
|
||
response['Access-Control-Allow-Origin'] = origin
|
||
response['Access-Control-Allow-Credentials'] = 'true'
|
||
|
||
return response
|
||
|
||
|
||
@api_view(['POST'])
|
||
@permission_classes([AllowAny])
|
||
def push_test_event(request):
|
||
"""
|
||
向所有连接的客户端广播一个 SSE 测试事件。
|
||
"""
|
||
services.push_simple_message_with_object_id('order_paid', '订单已支付', 12345)
|
||
return Response({
|
||
'status': 'ok',
|
||
'message': 'Test event broadcasted',
|
||
'clients': len(services.get_active_connections())
|
||
})
|
||
|
||
|
||
@api_view(['GET'])
|
||
@permission_classes([AllowAny])
|
||
def get_sse_status(request):
|
||
"""
|
||
获取 SSE 连接状态信息。
|
||
|
||
返回：
|
||
- clients: 当前连接的客户端数量
|
||
- status: 服务状态
|
||
"""
|
||
active_connections = services.get_active_connections()
|
||
return Response({
|
||
'status': 'running',
|
||
'clients': len(active_connections),
|
||
'message': f'SSE server is running with {len(active_connections)} active connection(s)',
|
||
})
|
||
|
||
|
||
@api_view(['POST'])
|
||
@permission_classes([AllowAny])
|
||
def shutdown_sse(request):
|
||
"""
|
||
优雅关闭所有SSE连接的端点
|
||
"""
|
||
services.push_sse_event_to_all({
|
||
'type': 'server_shutdown',
|
||
'message': 'Server is shutting down, please reconnect later'
|
||
})
|
||
|
||
return Response({
|
||
'status': 'ok',
|
||
'message': 'Shutdown signal sent to all SSE connections',
|
||
'clients': len(services.get_active_connections())
|
||
})
|