Coverage for src/lexigram/admin/services/realtime.py: 42%
130 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 14:56 +0800
1"""Real-time Data Service - WebSocket and SSE support for live updates."""
3from __future__ import annotations
5import asyncio
6from dataclasses import dataclass, field
7from datetime import datetime
8from enum import StrEnum
9from typing import TYPE_CHECKING, Any
11from lexigram.serialization import dumps_str, loads_str
13if TYPE_CHECKING:
14 from collections.abc import Awaitable, Callable
17class UpdateType(StrEnum):
18 """Types of real-time updates."""
20 METRIC = "metric"
21 CHART = "chart"
22 TABLE = "table"
23 NOTIFICATION = "notification"
24 CUSTOM = "custom"
27@dataclass
28class RealtimeUpdate:
29 """Real-time update message."""
31 type: UpdateType
32 data: dict[str, Any]
33 timestamp: datetime = field(default_factory=datetime.now)
34 target: str = "" # Target element ID
36 def to_json(self) -> str:
37 """Serialize to JSON."""
38 return dumps_str(
39 {
40 "type": self.type.value,
41 "data": self.data,
42 "timestamp": self.timestamp.isoformat(),
43 "target": self.target,
44 },
45 )
47 @classmethod
48 def from_json(cls, data: str) -> RealtimeUpdate:
49 """Deserialize from JSON."""
50 obj = loads_str(data)
51 return cls(
52 type=UpdateType(obj["type"]),
53 data=obj["data"],
54 timestamp=datetime.fromisoformat(obj["timestamp"]),
55 target=obj.get("target", ""),
56 )
59class ConnectionManager:
60 """Manages WebSocket connections."""
62 def __init__(self) -> None:
63 """Initialize connection manager."""
64 self.active_connections: dict[str, Any] = {}
65 self.subscriptions: dict[str, set[str]] = {} # topic -> set of connection_ids
67 async def connect(self, connection_id: str, websocket: Any) -> None:
68 """Connect a WebSocket client."""
69 self.active_connections[connection_id] = websocket
71 def disconnect(self, connection_id: str) -> None:
72 """Disconnect a WebSocket client."""
73 self.active_connections.pop(connection_id, None)
74 # Remove from all subscriptions
75 for topic_subs in self.subscriptions.values():
76 topic_subs.discard(connection_id)
78 def subscribe(self, connection_id: str, topic: str) -> None:
79 """Subscribe a connection to a topic."""
80 if topic not in self.subscriptions:
81 self.subscriptions[topic] = set()
82 self.subscriptions[topic].add(connection_id)
84 def unsubscribe(self, connection_id: str, topic: str) -> None:
85 """Unsubscribe a connection from a topic."""
86 if topic in self.subscriptions:
87 self.subscriptions[topic].discard(connection_id)
89 async def send_personal(self, connection_id: str, message: str) -> None:
90 """Send message to specific connection."""
91 websocket = self.active_connections.get(connection_id)
92 if websocket:
93 try:
94 await websocket.send_text(message)
95 except (ConnectionError, OSError, RuntimeError) as e:
96 from lexigram.logging import get_logger
98 logger = get_logger(__name__)
99 logger.debug(
100 "Error sending message to connection %s: %s",
101 connection_id,
102 e,
103 )
104 # Ensure the faulty connection is removed
105 self.disconnect(connection_id)
107 async def broadcast(self, message: str, topic: str | None = None) -> None:
108 """Broadcast message to all connections or specific topic."""
109 if topic and topic in self.subscriptions:
110 targets = self.subscriptions[topic]
111 else:
112 targets = set(self.active_connections.keys())
114 for connection_id in list(targets):
115 await self.send_personal(connection_id, message)
117 def get_connection_count(self, topic: str | None = None) -> int:
118 """Get number of active connections."""
119 if topic:
120 return len(self.subscriptions.get(topic, set()))
121 return len(self.active_connections)
124class RealtimeService:
125 """DI-injectable real-time data service."""
127 def __init__(self) -> None:
128 """Initialize real-time service."""
129 self.connection_manager = ConnectionManager()
130 self.update_handlers: dict[
131 str,
132 Callable[[RealtimeUpdate], Awaitable[None]],
133 ] = {}
135 async def connect(self, connection_id: str, websocket: Any) -> None:
136 """Connect a WebSocket client."""
137 await self.connection_manager.connect(connection_id, websocket)
139 def disconnect(self, connection_id: str) -> None:
140 """Disconnect a WebSocket client."""
141 self.connection_manager.disconnect(connection_id)
143 def subscribe(self, connection_id: str, topic: str) -> None:
144 """Subscribe connection to topic."""
145 self.connection_manager.subscribe(connection_id, topic)
147 def unsubscribe(self, connection_id: str, topic: str) -> None:
148 """Unsubscribe connection from topic."""
149 self.connection_manager.unsubscribe(connection_id, topic)
151 async def send_update(
152 self,
153 update: RealtimeUpdate,
154 connection_id: str | None = None,
155 topic: str | None = None,
156 ) -> None:
157 """
158 Send real-time update.
160 Args:
161 update: Update to send
162 connection_id: Send to specific connection (optional)
163 topic: Broadcast to topic (optional)
164 """
165 message = update.to_json()
167 if connection_id:
168 await self.connection_manager.send_personal(connection_id, message)
169 else:
170 await self.connection_manager.broadcast(message, topic)
172 async def send_metric_update(
173 self,
174 metric_name: str,
175 value: Any,
176 target: str = "",
177 topic: str | None = None,
178 ) -> None:
179 """Send metric update."""
180 update = RealtimeUpdate(
181 type=UpdateType.METRIC,
182 data={"name": metric_name, "value": value},
183 target=target,
184 )
185 await self.send_update(update, topic=topic)
187 async def send_chart_update(
188 self,
189 chart_data: dict[str, Any],
190 target: str = "",
191 topic: str | None = None,
192 ) -> None:
193 """Send chart data update."""
194 update = RealtimeUpdate(type=UpdateType.CHART, data=chart_data, target=target)
195 await self.send_update(update, topic=topic)
197 async def send_notification(
198 self,
199 message: str,
200 severity: str = "info",
201 topic: str | None = None,
202 ) -> None:
203 """Send notification."""
204 update = RealtimeUpdate(
205 type=UpdateType.NOTIFICATION,
206 data={"message": message, "severity": severity},
207 )
208 await self.send_update(update, topic=topic)
210 def register_handler(
211 self,
212 update_type: UpdateType,
213 handler: Callable[[RealtimeUpdate], Awaitable[None]],
214 ) -> None:
215 """Register update handler."""
216 self.update_handlers[update_type.value] = handler
218 async def handle_update(self, update: RealtimeUpdate) -> None:
219 """Handle incoming update."""
220 handler = self.update_handlers.get(update.type.value)
221 if handler:
222 await handler(update)
224 def get_stats(self) -> dict[str, Any]:
225 """Get service statistics."""
226 return {
227 "total_connections": self.connection_manager.get_connection_count(),
228 "topics": {
229 topic: self.connection_manager.get_connection_count(topic)
230 for topic in self.connection_manager.subscriptions
231 },
232 "handlers": list(self.update_handlers.keys()),
233 }
236class SSEEmitter:
237 """Server-Sent Events emitter as fallback."""
239 def __init__(self) -> None:
240 """Initialize SSE emitter."""
241 self.clients: dict[str, asyncio.Queue] = {}
243 def connect(self, client_id: str) -> asyncio.Queue:
244 """Connect SSE client."""
245 queue: asyncio.Queue = asyncio.Queue()
246 self.clients[client_id] = queue
247 return queue
249 def disconnect(self, client_id: str) -> None:
250 """Disconnect SSE client."""
251 self.clients.pop(client_id, None)
253 async def emit(
254 self,
255 event: str,
256 data: dict[str, Any],
257 client_id: str | None = None,
258 ) -> None:
259 """Emit SSE event."""
260 # Prepare message payload, logging failures to serialize
261 try:
262 payload = dumps_str(data)
263 except (TypeError, ValueError):
264 from lexigram.logging import get_logger
266 logger = get_logger(__name__)
267 logger.exception("Failed to JSON-serialize SSE data for event %s", event)
268 payload = str(data)
270 message = f"event: {event}\ndata: {payload}\n\n"
272 if client_id and client_id in self.clients:
273 try:
274 await self.clients[client_id].put(message)
275 except RuntimeError:
276 from lexigram.logging import get_logger
278 logger = get_logger(__name__)
279 logger.exception("Failed to put message on client queue %s", client_id)
280 else:
281 # Broadcast to all
282 for queue in list(self.clients.values()):
283 try:
284 await queue.put(message)
285 except (RuntimeError, OSError) as e:
286 from lexigram.logging import get_logger
288 logger = get_logger(__name__)
289 logger.debug("Failed to put message on broadcast queue: %s", e)
291 def get_client_count(self) -> int:
292 """Get number of connected clients."""
293 return len(self.clients)