86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
import json
|
|
import queue
|
|
import uuid
|
|
from typing import Dict, Any
|
|
import logging
|
|
|
|
from flask import Response
|
|
|
|
from events.types.event import Event
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Maps user_id → dict of {connection_id: queue}
|
|
user_queues: Dict[str, Dict[str, queue.Queue]] = {}
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
def get_queue(user_id: str, connection_id: str) -> queue.Queue:
|
|
"""Get or create a queue for a specific user connection."""
|
|
if user_id not in user_queues:
|
|
user_queues[user_id] = {}
|
|
|
|
if connection_id not in user_queues[user_id]:
|
|
user_queues[user_id][connection_id] = queue.Queue()
|
|
|
|
return user_queues[user_id][connection_id]
|
|
|
|
|
|
def send_to_user(user_id: str, data: Dict[str, Any]):
|
|
"""Send data to all connections for a specific user."""
|
|
logger.info(f"Sending data to {user_id} user quesues are {user_queues}")
|
|
if user_id in user_queues:
|
|
logger.info(f"Queued {user_id}")
|
|
# Format as SSE message once
|
|
message = f"data: {json.dumps(data)}\n\n".encode('utf-8')
|
|
|
|
# Send to all connections for this user
|
|
for connection_id, q in user_queues[user_id].items():
|
|
try:
|
|
logger.info(f"Sending message to {connection_id}")
|
|
q.put(message)
|
|
q.put(message, block=False)
|
|
except queue.Full:
|
|
# Skip if queue is full (connection might be dead)
|
|
pass
|
|
|
|
|
|
def send_event_to_user(user_id: str, event: Event):
|
|
"""Send an Event to all connections for a specific user."""
|
|
send_to_user(user_id, event.to_dict())
|
|
|
|
|
|
def sse_response_for_user(user_id: str):
|
|
"""Create SSE response for a user connection."""
|
|
# Generate unique connection ID
|
|
connection_id = str(uuid.uuid4())
|
|
user_queue = get_queue(user_id, connection_id)
|
|
logger.info(f"SSE response for {user_id} user is {user_queue}")
|
|
|
|
def generate():
|
|
try:
|
|
while True:
|
|
# Get message from queue (blocks until available)
|
|
message = user_queue.get()
|
|
logging.info(
|
|
f"Sending SSE message to user {user_id} (connection {connection_id}): {message.decode('utf-8').strip()}")
|
|
yield message
|
|
except GeneratorExit:
|
|
# Clean up when client disconnects
|
|
if user_id in user_queues and connection_id in user_queues[user_id]:
|
|
del user_queues[user_id][connection_id]
|
|
|
|
# Remove user entry if no connections remain
|
|
if not user_queues[user_id]:
|
|
del user_queues[user_id]
|
|
|
|
return Response(
|
|
generate(),
|
|
mimetype='text/event-stream',
|
|
headers={
|
|
'Cache-Control': 'no-cache',
|
|
'X-Accel-Buffering': 'no',
|
|
'Connection': 'keep-alive'
|
|
}
|
|
)
|