Coverage for src/lexigram/graphql/subscriptions/transport/_messaging.py: 50%

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._connection import SubscriptionConnection 

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 

27@dataclass 

28class _WSMessagingMixin: 

29 """Send/error/complete framing plus keepalive and cleanup.""" 

30 

31 _websocket: WebSocketProtocol | None 

32 _connection: SubscriptionConnection | None 

33 _connection_ack_sent: bool 

34 keepalive_interval: float 

35 

36 async def _send_message(self, message: dict[str, Any]) -> None: 

37 """Send a message to the client. 

38 

39 Args: 

40 message: The message to send. 

41 """ 

42 if not self._websocket: 

43 return 

44 try: 

45 await self._websocket.send_json(message) 

46 except Exception as e: 

47 if type(e).__name__ == "WebSocketDisconnect": 

48 logger.debug("WebSocket disconnected while sending: %s", e) 

49 else: 

50 raise 

51 

52 async def _send_next(self, subscription_id: str, data: Any) -> None: 

53 """Send next message (subscription data). 

54 

55 Args: 

56 subscription_id: The subscription ID. 

57 data: The data to send. 

58 """ 

59 # Convert ExecutionResult to dict if needed 

60 payload_data = None 

61 if data is not None: 

62 if hasattr(data, "data"): 

63 payload_data = data.data 

64 elif hasattr(data, "__dict__"): 

65 payload_data = data.__dict__ 

66 else: 

67 payload_data = data 

68 

69 await self._send_message( 

70 { 

71 "id": subscription_id, 

72 "type": GQLWSMessageType.NEXT, 

73 "payload": {"data": payload_data} if payload_data else {}, 

74 } 

75 ) 

76 

77 async def _send_error(self, subscription_id: str | None, error: str) -> None: 

78 """Send error message. 

79 

80 Args: 

81 subscription_id: The subscription ID (may be None for connection errors). 

82 error: The error message. 

83 """ 

84 await self._send_message( 

85 { 

86 "id": subscription_id, 

87 "type": GQLWSMessageType.ERROR, 

88 "payload": {"message": error}, 

89 } 

90 ) 

91 

92 async def _send_complete(self, subscription_id: str) -> None: 

93 """Send complete message. 

94 

95 Args: 

96 subscription_id: The subscription ID. 

97 """ 

98 await self._send_message( 

99 { 

100 "id": subscription_id, 

101 "type": GQLWSMessageType.COMPLETE, 

102 } 

103 ) 

104 

105 async def _keepalive(self) -> None: 

106 """Send keep-alive ping messages (graphql-transport-ws protocol).""" 

107 try: 

108 while True: 

109 await asyncio.sleep(self.keepalive_interval) 

110 if self._websocket and self._connection_ack_sent: 

111 await self._send_message({"type": GQLWSMessageType.PING}) 

112 except asyncio.CancelledError: 

113 pass 

114 

115 async def _cleanup(self) -> None: 

116 """Clean up resources on disconnect.""" 

117 if self._connection: 

118 # Cancel all active subscriptions 

119 for sub_id in list(self._connection.subscriptions.keys()): 

120 self._connection.remove(sub_id) 

121 self._connection = None 

122 self._websocket = None 

123 

124 # ------------------------------------------------------------------ 

125 # SubscriptionHandler protocol implementation 

126 # ------------------------------------------------------------------