Coverage for src/lexigram/graphql/subscriptions/transport/_transport.py: 36%
196 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 04:37 +0800
1"""GraphQL WebSocket Transport for subscriptions.
3Provides WebSocket transport implementation using graphql-transport-ws protocol.
4"""
6from __future__ import annotations
8import asyncio
9from dataclasses import dataclass, field
10from typing import TYPE_CHECKING, Any, cast
12from lexigram import serialization as json
13from lexigram.graphql.subscriptions.protocol import GQLWSMessageType
14from lexigram.graphql.subscriptions.transport._connection import SubscriptionConnection
15from lexigram.graphql.types import SubscriptionInfo
16from lexigram.logging import get_logger
18if TYPE_CHECKING:
19 from collections.abc import AsyncIterator, Awaitable, Callable
21 from lexigram.contracts.graphql.protocols import SubscriptionAuthHandlerProtocol
22 from lexigram.contracts.web import WebSocketProtocol
24logger = get_logger(__name__)
26from lexigram.graphql.subscriptions.transport._messaging import _WSMessagingMixin
29@dataclass
30class GraphQLWSTransport(_WSMessagingMixin):
31 """WebSocket transport for GraphQL subscriptions.
33 Implements the graphql-transport-ws protocol.
34 """
36 def __init__(
37 self,
38 execute: Callable[..., Awaitable[Any]] | None = None,
39 subscribe: Callable[..., Awaitable[Any]] | None = None,
40 connection_init_timeout: float = 10.0,
41 keepalive_interval: float = 30.0,
42 context_factory: Any | None = None,
43 auth_handler: Any | None = None,
44 subscription_auth_handler: SubscriptionAuthHandlerProtocol | None = None,
45 ):
46 """Initialize the transport.
48 Args:
49 execute: Function to execute GraphQL operations.
50 subscribe: Function to subscribe to GraphQL operations.
51 connection_init_timeout: Timeout for connection_init in seconds.
52 keepalive_interval: Interval for keep-alive messages in seconds.
53 context_factory: Optional factory to create GraphQL context.
54 auth_handler: Optional handler for connection authentication.
55 subscription_auth_handler: Optional per-subscription authorization
56 handler. When provided, :meth:`_handle_subscribe` calls
57 :meth:`~lexigram.contracts.graphql.protocols.SubscriptionAuthHandlerProtocol.authorize`
58 before any subscription setup. A ``False`` return value rejects
59 the subscription immediately.
60 """
61 self._execute = execute
62 self._subscribe = subscribe
63 self.connection_init_timeout = connection_init_timeout
64 self.keepalive_interval = keepalive_interval
65 self._context_factory = context_factory
66 self._auth_handler = auth_handler
67 self._subscription_auth_handler = subscription_auth_handler
68 self._connection: SubscriptionConnection | None = None
69 self._websocket: WebSocketProtocol | None = None
70 self._user: Any | None = None
71 self._connection_init_received = False
72 self._connection_ack_sent = False
73 self._background_tasks: set[asyncio.Task] = set()
75 async def handle(
76 self, websocket: WebSocketProtocol, app: Any | None = None
77 ) -> None:
78 """Handle a WebSocket connection.
80 Args:
81 websocket: The WebSocket connection; any object satisfying
82 :class:`~lexigram.contracts.http.WebSocketProtocol`
83 (e.g. Starlette's ``WebSocket``).
84 app: Optional application instance (currently unused; context_factory
85 and auth_handler must be injected via constructor).
86 """
87 self._websocket = websocket
88 self._connection = SubscriptionConnection()
90 await websocket.accept(subprotocol="graphql-transport-ws")
91 logger.debug("WebSocket accepted, starting receive loop")
93 # Start keep-alive task
94 keepalive_task = asyncio.create_task(self._keepalive())
96 try:
97 await self._receive_loop()
98 except (OSError, RuntimeError) as e:
99 logger.error("WebSocket error: %s", e)
100 import traceback
102 logger.error("Traceback: %s", traceback.format_exc())
103 except Exception as e:
104 if type(e).__name__ == "WebSocketDisconnect":
105 logger.debug("WebSocket disconnected during receive: %s", e)
106 else:
107 logger.error("WebSocket error: %s", e, exc_info=True)
108 finally:
109 keepalive_task.cancel()
110 await self._cleanup()
112 async def _receive_loop(self) -> None:
113 """Receive and process messages from the client."""
114 if not self._websocket:
115 return
117 logger.info("===== STARTING RECEIVE LOOP =====")
119 while True:
120 try:
121 data = await self._websocket.receive_text()
122 logger.info("WS RECEIVED: {data[:200]}")
123 except (OSError, RuntimeError) as e:
124 logger.debug("Receive exception: %s", e)
125 break
126 except Exception as e:
127 # WebSocketDisconnect extends Exception directly, not OSError/RuntimeError
128 if type(e).__name__ == "WebSocketDisconnect":
129 logger.debug("WebSocket disconnected: %s", e)
130 break
131 raise
133 try:
134 await self._handle_message(json.loads(data))
135 except json.JSONDecodeError as e:
136 await self._send_error(None, f"Invalid JSON: {e}")
137 except (RuntimeError, ValueError, TypeError) as e:
138 logger.error("Error handling message: %s", e)
139 await self._send_error(None, str(e))
141 async def _handle_message(self, message: dict[str, Any]) -> None:
142 """Handle a single message from the client.
144 Args:
145 message: The parsed message.
146 """
147 msg_type = message.get("type")
148 logger.info("WS message: %s", msg_type)
149 payload = message.get("payload", {})
150 subscription_id = message.get("id")
152 if msg_type == GQLWSMessageType.CONNECTION_INIT:
153 await self._handle_connection_init(payload)
154 elif msg_type == GQLWSMessageType.SUBSCRIBE:
155 await self._handle_subscribe(subscription_id, payload) # type: ignore[arg-type]
156 elif msg_type == GQLWSMessageType.COMPLETE:
157 await self._handle_complete(subscription_id) # type: ignore[arg-type]
158 elif msg_type == GQLWSMessageType.PING:
159 await self._send_message({"type": GQLWSMessageType.PONG})
160 elif msg_type == GQLWSMessageType.PONG:
161 pass # Client acknowledged our ping — nothing to do
162 elif not self._connection_ack_sent:
163 await self._send_error(
164 subscription_id,
165 f"Unexpected message before connection_ack: {msg_type}",
166 )
168 async def _handle_connection_init(self, payload: dict[str, Any]) -> None:
169 """Handle connection_init message.
171 Args:
172 payload: The connection payload.
173 """
174 self._connection_init_received = True
175 logger.debug("WS connection_init payload: %s", payload)
177 # Authentication via auth_handler if provided
178 if self._auth_handler:
179 try:
180 # authenticate() might return a user object or bool
181 auth_result = await self._auth_handler.authenticate(payload)
182 if auth_result and not isinstance(auth_result, bool):
183 self._user = auth_result
184 logger.debug("WS auth user set: %s", self._user)
185 elif auth_result is True:
186 # Success but no user object returned, check payload for user
187 self._user = payload.get("user")
188 logger.debug("WS auth using payload user: %s", self._user)
189 except (RuntimeError, ValueError, TypeError, LookupError) as e:
190 logger.warning("WebSocket auth failed: %s", e)
191 await self._send_message(
192 {
193 "type": GQLWSMessageType.CONNECTION_ERROR,
194 "payload": {"message": str(e)},
195 }
196 )
197 if self._websocket:
198 await self._websocket.close(code=4403)
199 return
201 await self._send_message({"type": GQLWSMessageType.CONNECTION_ACK})
202 self._connection_ack_sent = True
204 async def _handle_subscribe(
205 self,
206 subscription_id: str,
207 payload: dict[str, Any],
208 ) -> None:
209 """Handle subscribe message.
211 Args:
212 subscription_id: The subscription ID.
213 payload: The subscribe payload with query, variables, operationName.
214 """
215 logger.debug(
216 "_handle_subscribe called with subscription_id=%s, payload=%s",
217 subscription_id,
218 payload,
219 )
220 if not self._connection:
221 await self._send_error(subscription_id, "No connection")
222 return
224 query = payload.get("query")
225 variables = payload.get("variables")
226 operation_name = payload.get("operationName")
228 if not query:
229 await self._send_error(subscription_id, "Missing query")
230 return
232 # Per-subscription authorization check — runs after query validation so
233 # the auth handler receives the actual operation details.
234 if self._subscription_auth_handler is not None:
235 try:
236 allowed = await self._subscription_auth_handler.authorize(
237 user=self._user,
238 operation_name=operation_name,
239 query=query,
240 )
241 except (
242 RuntimeError,
243 ValueError,
244 TypeError,
245 LookupError,
246 PermissionError,
247 ) as e:
248 logger.warning("subscription_auth_error", error=str(e))
249 await self._send_error(subscription_id, "Authorization error")
250 return
252 if not allowed:
253 await self._send_error(subscription_id, "Unauthorized subscription")
254 return
256 # Store subscription info
257 info = SubscriptionInfo(
258 subscription_id=subscription_id,
259 operation_name=operation_name,
260 query=query,
261 variables=variables or {},
262 )
263 self._connection.add(subscription_id, info)
265 try:
266 # Check if this is a subscription query
267 is_subscription = query.strip().lower().startswith("subscription")
268 logger.debug(
269 "Handling subscribe, is_subscription: %s, query: %s...",
270 is_subscription,
271 query[:50],
272 )
274 # Build context if factory is available
275 context_value = None
276 logger.debug(
277 "Context factory available: %s", self._context_factory is not None
278 )
279 logger.debug("User to set in context: %s", self._user)
280 if self._context_factory:
281 from lexigram.graphql.core.context import GraphQLRequest as GQLRequest
283 req_obj = GQLRequest(
284 query=query,
285 variables=variables or {},
286 operation_name=operation_name,
287 )
288 # Correct arguments for ContextFactory.create_context
289 metadata: dict[str, Any] = {"raw_request": self._websocket}
290 context_value = await self._context_factory.create_context(
291 request=req_obj,
292 user=self._user,
293 metadata=metadata,
294 )
295 logger.debug("Created context with user: %s", self._user)
296 else:
297 # Create basic context with user even without factory (must be an object, not dict)
298 from lexigram.graphql.core.context import GraphQLContext
300 ws_metadata: dict[str, Any] = {"request": self._websocket}
301 context_value = GraphQLContext(
302 user=self._user,
303 metadata=ws_metadata,
304 )
305 logger.debug("Created basic context with user: %s", self._user)
307 if is_subscription:
308 # Execute the subscription
309 if self._subscribe:
310 result = await self._subscribe(
311 query,
312 variable_values=variables,
313 operation_name=operation_name,
314 context_value=context_value,
315 )
317 # Check if it's an iterable (subscription result)
318 if hasattr(result, "__aiter__"):
319 task = asyncio.create_task(
320 self._stream_subscription(subscription_id, result),
321 )
322 self._background_tasks.add(task)
323 task.add_done_callback(self._background_tasks.discard)
324 else:
325 # Immediate result (not a stream)
326 await self._send_next(subscription_id, result)
327 await self._send_complete(subscription_id)
328 else:
329 await self._send_error(
330 subscription_id,
331 "Subscription handler not configured",
332 )
333 elif self._execute:
334 # Execute query/mutation via WebSocket
335 result = await self._execute(
336 query,
337 variable_values=variables,
338 operation_name=operation_name,
339 context_value=context_value,
340 )
341 await self._send_next(subscription_id, result)
342 await self._send_complete(subscription_id)
343 else:
344 await self._send_error(
345 subscription_id,
346 "Execute handler not configured",
347 )
349 except (RuntimeError, ValueError, TypeError, LookupError) as e:
350 logger.error("Subscribe error: %s", e)
351 await self._send_error(subscription_id, str(e))
352 self._connection.remove(subscription_id)
354 async def _stream_subscription(
355 self,
356 subscription_id: str,
357 result: Any,
358 ) -> None:
359 """Stream subscription results to the client.
361 Args:
362 subscription_id: The subscription ID.
363 result: The subscription result iterator.
364 """
365 try:
366 async for item in result:
367 await self._send_next(subscription_id, item)
368 except (RuntimeError, ValueError, TypeError, LookupError) as e:
369 logger.error("Stream error: %s", e, exc_info=True)
370 try:
371 await self._send_error(subscription_id, str(e))
372 except (OSError, RuntimeError) as send_err:
373 logger.debug("subscription_send_error_failed", error=str(send_err))
374 finally:
375 try:
376 await self._send_complete(subscription_id)
377 except (OSError, RuntimeError) as complete_err:
378 logger.debug(
379 "subscription_send_complete_failed", error=str(complete_err)
380 )
381 if self._connection:
382 self._connection.remove(subscription_id)
384 async def _handle_complete(self, subscription_id: str) -> None:
385 """Handle complete message (client wants to end subscription).
387 Args:
388 subscription_id: The subscription ID.
389 """
390 if self._connection:
391 self._connection.remove(subscription_id)
393 async def subscribe(
394 self,
395 field_name: str,
396 args: dict[str, Any],
397 context: Any,
398 info: Any,
399 ) -> AsyncIterator[Any]:
400 """Satisfy the SubscriptionHandler protocol by building and executing a subscription.
402 Constructs a minimal GraphQL subscription document from *field_name*
403 and *args*, delegates to the underlying ``_subscribe`` callable, and
404 returns an async iterator of result values. Single (non-iterator)
405 return values are wrapped in a one-element async iterator for
406 protocol uniformity.
408 Args:
409 field_name: The GraphQL subscription field to subscribe to.
410 args: Keyword arguments forwarded as query variables and
411 injected into the subscription field call-site.
412 context: GraphQL context value passed to the executor.
413 info: ResolverProtocol info (ignored at transport level but required
414 by the SubscriptionHandler protocol).
416 Returns:
417 An async iterator that yields subscription event values.
419 Raises:
420 RuntimeError: If no ``_subscribe`` callable has been configured.
421 """
422 if self._subscribe is None:
423 raise RuntimeError("No subscribe callable configured on GraphQLWSTransport")
424 if args:
425 var_defs = ", ".join(f"${k}: String" for k in args)
426 arg_ref = ", ".join(f"{k}: ${k}" for k in args)
427 query = f"subscription Sub({var_defs}) {{ {field_name}({arg_ref}) }}"
428 else:
429 query = f"subscription {{ {field_name} }}"
431 result = await self._subscribe(
432 query,
433 variable_values=args or None,
434 context_value=context,
435 )
437 if hasattr(result, "__aiter__"):
438 return cast("AsyncIterator[Any]", result)
440 async def _single_event() -> AsyncIterator[Any]:
441 yield result
443 return _single_event()