feat: big version, added tasks for backup_database and stock change, added health check api, approve sse (support channel via merchant)

2025-11-26 21:49:42 +08:00
parent 6bf0465d05
commit a9c75a13fa
26 changed files with 7158 additions and 216 deletions
--- a/sse/services.py
+++ b/sse/services.py
@@ -3,50 +3,81 @@ import logging

 logger = logging.getLogger(__name__)

-# 存储所有活动的 SSE 连接队列（同步队列）
-_connections = set()
+# 存储所有活动的 SSE 连接队列（按商户ID组织）
+_connections = {}  # {merchant_id: set([conn_queue1, conn_queue2, ...])}


-def get_active_connections():
-    """获取当前所有活动的 SSE 连接队列"""
-    return _connections
+def get_all_connections_count():
+    """获取所有商户的连接总数"""
+    total = 0
+    for merchant_id, connections in _connections.items():
+        total += len(connections)
+    return total


-def push_connection(conn_queue):
+def get_merchant_connections(merchant_id):
+    """获取指定商户的所有连接队列"""
+    return _connections.get(merchant_id, set())
+
+
+def get_all_merchant_ids():
+    """获取所有有活跃连接的商户ID列表"""
+    return list(_connections.keys())
+
+
+def push_connection(merchant_id, conn_queue):
    """
-    将一个新的连接队列添加到活动连接集合中
+    将一个新的连接队列添加到指定商户的连接集合中
    
    参数:
+    - merchant_id: 商户ID
    - conn_queue: queue.Queue 实例
    """
-    _connections.add(conn_queue)
-    logger.info(f"新的 SSE 连接建立，当前连接数: {len(_connections)}")
+    if merchant_id not in _connections:
+        _connections[merchant_id] = set()
+    
+    _connections[merchant_id].add(conn_queue)
+    total_connections = sum(len(conns) for conns in _connections.values())
+    logger.info(f"商户 {merchant_id} 的新 SSE 连接建立，该商户连接数: {len(_connections[merchant_id])}，总连接数: {total_connections}")


-def remove_connection(conn_queue):
+def remove_connection(merchant_id, conn_queue):
    """
-    从活动连接集合中移除一个连接队列
+    从指定商户的连接集合中移除一个连接队列
    
    参数:
+    - merchant_id: 商户ID
    - conn_queue: queue.Queue 实例
    """
-    _connections.discard(conn_queue)
-    logger.info(f"SSE 连接断开，当前连接数: {len(_connections)}")
+    if merchant_id in _connections:
+        _connections[merchant_id].discard(conn_queue)
+        
+        # 如果该商户没有其他连接，则移除整个商户记录
+        if not _connections[merchant_id]:
+            del _connections[merchant_id]
+        
+        total_connections = sum(len(conns) for conns in _connections.values())
+        logger.info(f"商户 {merchant_id} 的 SSE 连接断开，该商户连接数: {len(_connections.get(merchant_id, set()))}，总连接数: {total_connections}")
+    else:
+        logger.warning(f"尝试从未注册的商户 {merchant_id} 移除连接")


 def cleanup_all_connections():
    """
    清理所有连接（用于服务器关闭时）
    """
-    connections_list = list(_connections)
-    for conn_queue in connections_list:
-        try:
-            # 发送关闭信号到队列
-            if hasattr(conn_queue, 'put_nowait'):
-                conn_queue.put_nowait({'type': 'server_shutdown', 'message': 'Server shutting down'})
-                logger.debug(f'shutdown connection {conn_queue}')
-        except Exception:
-            pass  # 忽略错误，因为连接可能已经断开
+    all_connections = {}
+    for merchant_id, connections in list(_connections.items()):
+        all_connections[merchant_id] = list(connections)
+        for conn_queue in connections:
+            try:
+                # 发送关闭信号到队列
+                if hasattr(conn_queue, 'put_nowait'):
+                    conn_queue.put_nowait({'type': 'server_shutdown', 'message': 'Server shutting down'})
+                    logger.debug(f'shutdown connection {conn_queue} for merchant {merchant_id}')
+            except Exception:
+                pass  # 忽略错误，因为连接可能已经断开
+    
    _connections.clear()
    logger.info("所有SSE连接已清理")

@@ -58,23 +89,60 @@ def push_sse_event_to_all(event_data: dict):
    参数:
    - event_data: 要发送的事件数据（字典）
    """
+    disconnected = {}
+    
+    for merchant_id, connections in list(_connections.items()):
+        disconnected[merchant_id] = []
+        
+        for conn_queue in connections:
+            try:
+                # 使用put_nowait非阻塞发送消息
+                conn_queue.put_nowait(event_data)
+            except queue.Full:
+                # 队列满了，说明客户端处理消息太慢，跳过这条消息
+                logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
+            except Exception as e:
+                # 其他异常可能表示连接已断开
+                logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
+                disconnected[merchant_id].append(conn_queue)
+    
+    # 清理断开的连接
+    for merchant_id, conn_queues in disconnected.items():
+        for conn_queue in conn_queues:
+            remove_connection(merchant_id, conn_queue)
+
+
+def push_sse_event_to_merchant(merchant_id, event_data: dict):
+    """
+    向指定商户的所有连接客户端广播一个 SSE 事件（同步队列版本）
+    
+    参数:
+    - merchant_id: 目标商户ID
+    - event_data: 要发送的事件数据（字典）
+    """
+    if merchant_id not in _connections:
+        logger.info(f"商户 {merchant_id} 没有活跃连接，跳过消息发送")
+        return
+    
    disconnected = []
    
-    for conn_queue in list(_connections):
+    for conn_queue in list(_connections[merchant_id]):
        try:
            # 使用put_nowait非阻塞发送消息
            conn_queue.put_nowait(event_data)
        except queue.Full:
            # 队列满了，说明客户端处理消息太慢，跳过这条消息
-            logger.warning(f"队列已满，跳过消息: {event_data.get('type', 'unknown')}")
+            logger.warning(f"商户 {merchant_id} 队列已满，跳过消息: {event_data.get('type', 'unknown')}")
        except Exception as e:
            # 其他异常可能表示连接已断开
-            logger.error(f"向队列发送消息失败: {e}")
+            logger.error(f"向商户 {merchant_id} 的队列发送消息失败: {e}")
            disconnected.append(conn_queue)
    
    # 清理断开的连接
    for conn_queue in disconnected:
-        remove_connection(conn_queue)
+        remove_connection(merchant_id, conn_queue)
+    
+    logger.info(f"向商户 {merchant_id} 广播 SSE 事件: type={event_data.get('type', 'unknown')}, 接收客户端数={len(_connections.get(merchant_id, set()))}")


 def push_simple_message_with_object_id(event_type: str, message: str, object_id):
@@ -93,4 +161,25 @@ def push_simple_message_with_object_id(event_type: str, message: str, object_id)
        'object_id': object_id,
    }
    push_sse_event_to_all(event_data)
-    logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={len(_connections)}")
+    total_connections = sum(len(conns) for conns in _connections.values())
+    logger.info(f"广播 SSE 事件: type={event_type}, object_id={object_id}, 接收客户端数={total_connections}")
+
+
+def push_simple_message_with_object_id_to_merchant(merchant_id, event_type: str, message: str, object_id):
+    """
+    向指定商户的所有连接客户端广播一个简单消息事件，包含关联对象ID
+    
+    参数:
+    - merchant_id: 目标商户ID
+    - event_type: 事件类型字符串
+    - message: 消息内容字符串
+    - object_id: 关联对象的ID（整数或字符串）
+    """
+    event_data = {
+        'mode': 'simple_message',
+        'type': event_type,
+        'message': message,
+        'object_id': object_id,
+        'merchant_id': merchant_id,  # 添加商户ID，便于客户端验证
+    }
+    push_sse_event_to_merchant(merchant_id, event_data)