Coverage for src/lexigram/graphql/core/batching.py: 94%

1"""GraphQL query batching support. 

2 

3This module provides batch execution of multiple GraphQL operations 

4in a single HTTP request. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from dataclasses import dataclass, field 

11from typing import Any, cast 

12 

13from lexigram.graphql.core.context import ( 

14 GraphQLContext, 

15 GraphQLRequest, 

16 GraphQLResponse, 

17) 

18from lexigram.graphql.exceptions import InputGraphQLError 

19from lexigram.logging import get_logger 

20 

21logger = get_logger(__name__) 

22 

23 

24@dataclass 

25class BatchResult: 

26 """Result of batch execution. 

27 

28 Attributes: 

29 responses: List of GraphQL responses. 

30 total_count: Total number of operations. 

31 success_count: Number of successful operations. 

32 error_count: Number of failed operations. 

33 """ 

34 

35 responses: list[GraphQLResponse[Any] | Exception] = field(default_factory=list) 

36 total_count: int = 0 

37 success_count: int = 0 

38 error_count: int = 0 

39 

40 

41class BatchExecutor: 

42 """Execute multiple GraphQL operations in a single request. 

43 

44 Allows clients to send an array of operations in a single HTTP request, 

45 improving performance for multiple independent queries. 

46 

47 Example: 

48 ```python 

49 executor = GraphQLExecutorProtocol(schema) 

50 batch_executor = BatchExecutor(executor, max_batch_size=10) 

51 

52 operations = [ 

53 GraphQLRequest(query="{ user(id: 1) { name } }"), 

54 GraphQLRequest(query="{ user(id: 2) { name } }"), 

55 ] 

56 

57 results = await batch_executor.execute_batch(operations, context) 

58 ``` 

59 """ 

60 

61 def __init__( 

62 self, 

63 executor: Any, # GraphQLExecutorProtocol 

64 max_batch_size: int = 10, 

65 ): 

66 """Initialize the batch executor. 

67 

68 Args: 

69 executor: The GraphQL executor to use. 

70 max_batch_size: Maximum number of operations per batch. 

71 """ 

72 self._executor = executor 

73 self._max_batch_size = max_batch_size 

74 

75 @property 

76 def max_batch_size(self) -> int: 

77 """Get maximum batch size.""" 

78 return self._max_batch_size 

79 

80 async def execute_batch( 

81 self, 

82 operations: list[GraphQLRequest], 

83 context: GraphQLContext, 

84 ) -> BatchResult: 

85 """Execute operations as a batch. 

86 

87 Args: 

88 operations: List of GraphQL requests to execute. 

89 context: GraphQL context. 

90 

91 Returns: 

92 BatchResult with all responses. 

93 

94 Raises: 

95 InputGraphQLError: If batch size exceeds limit. 

96 """ 

97 if len(operations) > self._max_batch_size: 

98 raise InputGraphQLError( 

99 message=f"Batch size {len(operations)} exceeds maximum of {self._max_batch_size}", 

100 ) 

101 

102 if not operations: 

103 return BatchResult() 

104 

105 # Create tasks for concurrent execution 

106 tasks = [] 

107 for op in operations: 

108 # Create a copy of context for each operation 

109 op_context = context.copy() if hasattr(context, "copy") else context 

110 task = self._execute_single(op, op_context) 

111 tasks.append(task) 

112 

113 # Execute all operations concurrently 

114 responses = await asyncio.gather(*tasks, return_exceptions=True) 

115 

116 # Count successes and errors 

117 success_count = sum(1 for r in responses if not isinstance(r, Exception)) 

118 error_count = len(responses) - success_count 

119 

120 return BatchResult( 

121 responses=cast("list[GraphQLResponse[Any] | Exception]", responses), 

122 total_count=len(operations), 

123 success_count=success_count, 

124 error_count=error_count, 

125 ) 

126 

127 async def _execute_single( 

128 self, 

129 operation: GraphQLRequest, 

130 context: GraphQLContext, 

131 ) -> GraphQLResponse[Any] | Exception: 

132 """Execute a single operation. 

133 

134 Args: 

135 operation: GraphQL request. 

136 context: GraphQL context. 

137 

138 Returns: 

139 GraphQL response or exception. 

140 """ 

141 try: 

142 # Create a new execution context for this operation 

143 from lexigram.graphql.core.execution import ExecutionContextProtocol 

144 

145 exec_context = ExecutionContextProtocol(context=context) 

146 return cast( 

147 "GraphQLResponse[Any]", 

148 await self._executor._execute_internal(exec_context), 

149 ) 

150 except Exception as e: # noqa: BLE001 — batch operations must not raise; exception is returned as the batch result 

151 logger.exception("Batch operation failed") 

152 return e 

153 

154 

155def is_batch_request(body: Any) -> bool: 

156 """Detect if a request body is a batch operation. 

157 

158 Args: 

159 body: The request body (parsed JSON). 

160 

161 Returns: 

162 True if the body is a list of operations. 

163 """ 

164 return isinstance(body, list) 

165 

166 

167def parse_batch_request(body: Any) -> list[GraphQLRequest]: 

168 """Parse a batch request body into GraphQLRequests. 

169 

170 Args: 

171 body: The request body (list of operations). 

172 

173 Returns: 

174 List of GraphQLRequest objects. 

175 """ 

176 requests = [] 

177 for item in body: 

178 if isinstance(item, dict): 

179 requests.append( 

180 GraphQLRequest( 

181 query=item.get("query"), 

182 variables=item.get("variables"), 

183 operation_name=item.get("operationName"), 

184 ) 

185 ) 

186 return requests 

187 

188 

189__all__ = [ 

190 "BatchExecutor", 

191 "BatchResult", 

192 "is_batch_request", 

193 "parse_batch_request", 

194]