Coverage for src/lexigram/graphql/subscriptions/transport/handler.py: 55%

1"""GraphQL WebSocket Transport for subscriptions. 

2 

3Provides WebSocket transport implementation using graphql-transport-ws protocol. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9from dataclasses import dataclass, field 

10from typing import TYPE_CHECKING, Any, cast 

11 

12from lexigram import serialization as json 

13from lexigram.graphql.subscriptions.protocol import GQLWSMessageType 

14from lexigram.graphql.subscriptions.transport._transport import GraphQLWSTransport 

15from lexigram.graphql.types import SubscriptionInfo 

16from lexigram.logging import get_logger 

17 

18if TYPE_CHECKING: 

19 from collections.abc import AsyncIterator, Awaitable, Callable 

20 

21 from lexigram.contracts.graphql.protocols import SubscriptionAuthHandlerProtocol 

22 from lexigram.contracts.web import WebSocketProtocol 

23 

24logger = get_logger(__name__) 

25 

26 

27async def graphql_ws_endpoint(websocket: Any) -> None: 

28 """ASGI-compatible endpoint for GraphQL WebSocket connections. 

29 

30 This is a simple endpoint that can be used with Starlette. 

31 For production use, configure with proper execute/subscribe handlers. 

32 

33 Example: 

34 from starlette.routing import Route 

35 from lexigram.graphql import constants as const 

36 

37 async def get_transport(): 

38 from lexigram.graphql.schema import build_schema 

39 # Configure with your schema 

40 return GraphQLWSTransport( 

41 execute=execute_fn, 

42 subscribe=subscribe_fn, 

43 ) 

44 

45 routes = [ 

46 Route(const.DEFAULT_SUBSCRIPTIONS_PATH, graphql_ws_endpoint), 

47 ] 

48 """ 

49 transport = GraphQLWSTransport() 

50 await transport.handle(websocket) 

51 

52 

53def create_ws_route(path: str) -> Any: 

54 """Create a WebSocket route for GraphQL subscriptions. 

55 

56 Returns a Starlette ``Route`` when ``starlette`` is installed. The return 

57 type is ``Any`` so callers are not forced to depend on Starlette directly. 

58 

59 Args: 

60 path: The WebSocket path. 

61 

62 Returns: 

63 A Starlette Route object. 

64 """ 

65 from starlette.routing import Route # deferred — Starlette is optional 

66 

67 return Route(path, graphql_ws_endpoint) 

68 

69 

70class GraphQLWSHandler: 

71 """Handler class for GraphQL WebSocket subscriptions. 

72 

73 This class provides a more configurable way to handle subscriptions. 

74 """ 

75 

76 def __init__( 

77 self, 

78 execute: Callable[..., Awaitable[Any]] | None = None, 

79 subscribe: Callable[..., Awaitable[Any]] | None = None, 

80 connection_init_timeout: float = 10.0, 

81 keepalive_interval: float = 30.0, 

82 context_factory: Any | None = None, 

83 auth_handler: Any | None = None, 

84 ): 

85 """Initialize the handler. 

86 

87 Args: 

88 execute: Function to execute GraphQL operations. 

89 subscribe: Function to subscribe to GraphQL operations. 

90 connection_init_timeout: Timeout for connection_init in seconds. 

91 keepalive_interval: Interval for keep-alive messages in seconds. 

92 context_factory: Optional factory to create GraphQL context. 

93 auth_handler: Optional handler for connection authentication. 

94 """ 

95 self._transport = GraphQLWSTransport( 

96 execute=execute, 

97 subscribe=subscribe, 

98 connection_init_timeout=connection_init_timeout, 

99 keepalive_interval=keepalive_interval, 

100 context_factory=context_factory, 

101 auth_handler=auth_handler, 

102 ) 

103 

104 async def __call__(self, scope: Any, receive: Any, send: Any) -> None: 

105 """Handle WebSocket connection. 

106 

107 Args: 

108 scope: ASGI scope 

109 receive: ASGI receive 

110 send: ASGI send 

111 """ 

112 from starlette.websockets import WebSocket 

113 

114 websocket = WebSocket(scope=scope, receive=receive, send=send) 

115 await self._transport.handle(websocket, app=scope.get("app")) 

116 

117 @classmethod 

118 def create_from_schema( 

119 cls, 

120 schema: Any, 

121 execute: Callable[..., Awaitable[Any]] | None = None, 

122 context_factory: Any | None = None, 

123 auth_handler: Any | None = None, 

124 ) -> GraphQLWSHandler: 

125 """Create a handler from a Strawberry schema. 

126 

127 Args: 

128 schema: The Strawberry schema. 

129 execute: Optional custom execute function. 

130 context_factory: Optional factory to create GraphQL context. 

131 auth_handler: Optional handler for connection authentication. 

132 

133 Returns: 

134 A configured GraphQLWSHandler. 

135 """ 

136 # Get execute/subscribe from schema if not provided 

137 if execute is None and hasattr(schema, "execute"): 

138 execute = schema.execute 

139 

140 subscribe = None 

141 if hasattr(schema, "subscribe"): 

142 subscribe = schema.subscribe 

143 

144 return cls( 

145 execute=execute, 

146 subscribe=subscribe, 

147 context_factory=context_factory, 

148 auth_handler=auth_handler, 

149 )