forked from erp-dev/erp
167 lines
5.9 KiB
Python
167 lines
5.9 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 IsAuthenticated
|
||
from rest_framework.response import Response
|
||
from rest_framework import status
|
||
from . import services
|
||
from .auth_utils import require_sse_authentication, get_user_merchant_id
|
||
import queue
|
||
import json
|
||
|
||
|
||
@csrf_exempt
|
||
@require_http_methods(["GET", "OPTIONS"])
|
||
@require_sse_authentication
|
||
def create_sse_event(request):
|
||
"""
|
||
创建一个 SSE 事件流响应。
|
||
客户端连接到此端点后会保持长连接，等待服务器推送事件。
|
||
|
||
使用方法：
|
||
- GET /sse/
|
||
- 保持连接打开以接收实时事件
|
||
|
||
注意：
|
||
- 使用纯 Django 视图，不使用 DRF，避免内容协商导致的 406 错误
|
||
- SSE 需要特殊的 CORS 配置
|
||
- 需要JWT认证且用户必须有关联的商户
|
||
"""
|
||
# 处理 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():
|
||
# 获取当前商户ID
|
||
merchant_id = request.merchant_id
|
||
|
||
# 创建一个同步队列用于接收消息
|
||
conn_queue = queue.Queue(maxsize=100)
|
||
services.push_connection(merchant_id, conn_queue)
|
||
|
||
try:
|
||
# 发送初始连接成功消息
|
||
yield f"data: {json.dumps({'type': 'connected', 'message': 'SSE connection established', 'merchant_id': merchant_id})}\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(merchant_id, 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([IsAuthenticated])
|
||
def push_test_event(request):
|
||
"""
|
||
向当前用户所属商户的所有连接客户端广播一个 SSE 测试事件。
|
||
"""
|
||
# 获取当前用户的商户ID
|
||
merchant_id = get_user_merchant_id(request)
|
||
if not merchant_id:
|
||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||
|
||
services.push_simple_message_with_object_id_to_merchant(
|
||
merchant_id, 'order_paid', '订单已支付', 12345
|
||
)
|
||
|
||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||
return Response({
|
||
'status': 'ok',
|
||
'message': 'Test event broadcasted to your merchant',
|
||
'merchant_id': merchant_id,
|
||
'clients': len(merchant_connections)
|
||
})
|
||
|
||
|
||
@api_view(['GET'])
|
||
@permission_classes([IsAuthenticated])
|
||
def get_sse_status(request):
|
||
"""
|
||
获取 SSE 连接状态信息。
|
||
|
||
返回：
|
||
- total_clients: 所有商户的客户端总数
|
||
- merchant_clients: 当前商户的客户端数量
|
||
- status: 服务状态
|
||
"""
|
||
merchant_id = get_user_merchant_id(request)
|
||
if not merchant_id:
|
||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||
|
||
# 获取所有连接数
|
||
all_connections = services.get_all_connections_count()
|
||
|
||
# 获取当前商户的连接数
|
||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||
|
||
return Response({
|
||
'status': 'running',
|
||
'total_clients': all_connections,
|
||
'merchant_clients': len(merchant_connections),
|
||
'merchant_id': merchant_id,
|
||
'message': f'SSE server is running with {all_connections} total connections, {len(merchant_connections)} for your merchant',
|
||
})
|
||
|
||
|
||
@api_view(['POST'])
|
||
@permission_classes([IsAuthenticated])
|
||
def shutdown_sse(request):
|
||
"""
|
||
关闭当前用户所属商户的所有SSE连接
|
||
"""
|
||
merchant_id = get_user_merchant_id(request)
|
||
if not merchant_id:
|
||
return Response({'error': 'User has no associated merchant'}, status=status.HTTP_403_FORBIDDEN)
|
||
|
||
# 获取连接数
|
||
merchant_connections = services.get_merchant_connections(merchant_id)
|
||
client_count = len(merchant_connections)
|
||
|
||
# 实际关闭连接
|
||
services.shutdown_merchant_connections(merchant_id)
|
||
|
||
return Response({
|
||
'status': 'ok',
|
||
'message': 'Your merchant\'s SSE connections have been closed',
|
||
'merchant_id': merchant_id,
|
||
'clients': client_count
|
||
})
|