Coverage for src/lexigram/graphql/monitoring/tracing.py: 35%

113 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 04:37 +0800

1"""GraphQL query tracing. 

2 

3This module provides tracing support for GraphQL queries, 

4including resolver timing and execution path tracking. 

5""" 

6 

7from __future__ import annotations 

8 

9from collections.abc import Awaitable, Callable, Iterator 

10from dataclasses import dataclass, field 

11from datetime import UTC, datetime 

12import functools 

13import time 

14from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast 

15 

16from strawberry.extensions import SchemaExtension 

17 

18from lexigram.logging import get_logger 

19 

20if TYPE_CHECKING: 

21 from lexigram.contracts.observability.tracing import SpanProtocol, TracerProtocol 

22 

23logger = get_logger(__name__) 

24 

25P = ParamSpec("P") 

26T = TypeVar("T") 

27 

28 

29@dataclass 

30class TraceSpan: 

31 """A span in the execution trace. 

32 

33 Attributes: 

34 name: Span name. 

35 start_time: Start timestamp. 

36 end_time: End timestamp. 

37 duration_ms: Duration in milliseconds. 

38 metadata: Additional span metadata. 

39 children: Child spans. 

40 """ 

41 

42 name: str 

43 start_time: datetime = field(default_factory=lambda: datetime.now(UTC)) 

44 end_time: datetime | None = None 

45 duration_ms: float = 0.0 

46 metadata: dict[str, Any] = field(default_factory=dict) 

47 children: list[TraceSpan] = field(default_factory=list) 

48 

49 def end(self) -> None: 

50 """End the span and calculate duration.""" 

51 self.end_time = datetime.now(UTC) 

52 self.duration_ms = (self.end_time - self.start_time).total_seconds() * 1000 

53 

54 def add_child(self, child: TraceSpan) -> None: 

55 """Add a child span.""" 

56 self.children.append(child) 

57 

58 def to_dict(self) -> dict[str, Any]: 

59 """Convert to dictionary. 

60 

61 Returns: 

62 Dictionary representation. 

63 """ 

64 return { 

65 "name": self.name, 

66 "start_time": self.start_time.isoformat(), 

67 "end_time": self.end_time.isoformat() if self.end_time else None, 

68 "duration_ms": round(self.duration_ms, 3), 

69 "metadata": self.metadata, 

70 "children": [c.to_dict() for c in self.children], 

71 } 

72 

73 

74@dataclass 

75class ExecutionTrace: 

76 """Complete execution trace. 

77 

78 Attributes: 

79 operation_name: Name of the operation. 

80 start_time: Execution start time. 

81 end_time: Execution end time. 

82 total_duration_ms: Total duration. 

83 root_span: Root span of the trace. 

84 resolver_count: Number of resolvers called. 

85 """ 

86 

87 operation_name: str | None = None 

88 start_time: datetime = field(default_factory=lambda: datetime.now(UTC)) 

89 end_time: datetime | None = None 

90 total_duration_ms: float = 0.0 

91 root_span: TraceSpan | None = None 

92 resolver_count: int = 0 

93 

94 def end(self) -> None: 

95 """End the trace.""" 

96 self.end_time = datetime.now(UTC) 

97 self.total_duration_ms = ( 

98 self.end_time - self.start_time 

99 ).total_seconds() * 1000 

100 

101 def to_dict(self) -> dict[str, Any]: 

102 """Convert to dictionary. 

103 

104 Returns: 

105 Dictionary representation (Apollo tracing format). 

106 """ 

107 return { 

108 "version": 1, 

109 "startTime": self.start_time.isoformat(), 

110 "endTime": self.end_time.isoformat() if self.end_time else None, 

111 "duration": int(self.total_duration_ms * 1_000_000), # nanoseconds 

112 "execution": { 

113 "resolvers": self.root_span.to_dict() if self.root_span else {}, 

114 }, 

115 } 

116 

117 

118class TracingExtension(SchemaExtension): 

119 """Strawberry extension for query tracing. 

120 

121 Adds tracing information to query responses in Apollo 

122 tracing format. 

123 

124 Example: 

125 ```python 

126 from lexigram.graphql.monitoring import TracingExtension 

127 

128 schema = strawberry.Schema( 

129 query=Query, 

130 extensions=[TracingExtension()], 

131 ) 

132 ``` 

133 """ 

134 

135 def __init__( 

136 self, 

137 include_in_response: bool = True, 

138 tracer: TracerProtocol | None = None, 

139 ) -> None: 

140 """Initialize the extension. 

141 

142 Args: 

143 include_in_response: Include tracing in response extensions. 

144 tracer: Optional kernel TracerProtocol for unified distributed tracing. 

145 When provided, GraphQL operation and resolver spans are forwarded 

146 to the application-wide tracing pipeline so GraphQL telemetry is 

147 correlated with infra traces. 

148 """ 

149 self._include_in_response = include_in_response 

150 self._tracer: TracerProtocol | None = tracer 

151 self._trace: ExecutionTrace | None = None 

152 self._current_span: TraceSpan | None = None 

153 self._kernel_root_span: SpanProtocol | None = None 

154 

155 def on_operation(self) -> Iterator[None]: 

156 """Hook called during operation execution.""" 

157 execution_context = self.execution_context 

158 

159 # Start kernel-level span for unified distributed tracing 

160 if self._tracer is not None: 

161 self._kernel_root_span = self._tracer.start_span( 

162 "graphql.operation", 

163 attributes={ 

164 "operation_name": execution_context.operation_name or "", 

165 "graphql.type": "operation", 

166 }, 

167 ) 

168 

169 # Start trace 

170 self._trace = ExecutionTrace( 

171 operation_name=execution_context.operation_name, 

172 ) 

173 self._trace.root_span = TraceSpan(name="operation") 

174 self._current_span = self._trace.root_span 

175 

176 yield 

177 

178 # End trace 

179 if self._trace: 

180 self._trace.root_span.end() 

181 self._trace.end() 

182 

183 # Add to response extensions 

184 if self._include_in_response and execution_context.result: 

185 if not execution_context.result.extensions: 

186 execution_context.result.extensions = {} 

187 execution_context.result.extensions["tracing"] = self._trace.to_dict() 

188 

189 # Close kernel span 

190 if self._kernel_root_span is not None: 

191 has_errors = bool( 

192 execution_context.result and execution_context.result.errors 

193 ) 

194 self._kernel_root_span.set_status("ERROR" if has_errors else "OK") 

195 self._kernel_root_span = None 

196 

197 def resolve( 

198 self, 

199 _next: Callable[..., Any], 

200 root: Any, 

201 info: Any, 

202 *args: Any, 

203 **kwargs: Any, 

204 ) -> Any: 

205 """Hook called during field resolution.""" 

206 if self._trace: 

207 self._trace.resolver_count += 1 

208 

209 # Create span for resolver 

210 field_name = info.field_name if hasattr(info, "field_name") else "unknown" 

211 parent_type: str | None = ( 

212 info.parent_type.name 

213 if hasattr(info, "parent_type") and info.parent_type 

214 else None 

215 ) 

216 return_type: str | None = ( 

217 str(info.return_type) if hasattr(info, "return_type") else None 

218 ) 

219 span = TraceSpan( 

220 name=field_name, 

221 metadata={ 

222 "parentType": parent_type, 

223 "returnType": return_type, 

224 }, 

225 ) 

226 

227 # Add as child of current span 

228 if self._current_span: 

229 self._current_span.add_child(span) 

230 

231 # Forward to kernel tracer for distributed tracing correlation 

232 kernel_span: SpanProtocol | None = None 

233 if self._tracer is not None: 

234 kernel_span = self._tracer.start_span( 

235 f"graphql.resolve.{field_name}", 

236 attributes={ 

237 "field_name": field_name, 

238 "parent_type": parent_type or "", 

239 "return_type": return_type or "", 

240 }, 

241 ) 

242 

243 try: 

244 return _next(root, info, *args, **kwargs) 

245 except Exception as _trace_err: # noqa: BLE001 — tracing wrapper must capture any error to mark the span before re-raising 

246 if kernel_span is not None: 

247 kernel_span.set_status("ERROR") 

248 raise 

249 finally: 

250 span.end() 

251 if kernel_span is not None: 

252 kernel_span.set_status("OK") 

253 

254 

255def trace_resolver( 

256 name: str | None = None, 

257) -> Callable[[Callable[P, T]], Callable[P, T]]: 

258 """Decorator for tracing resolver execution. 

259 

260 Args: 

261 name: Optional span name. 

262 

263 Returns: 

264 Decorator function. 

265 

266 Example: 

267 ```python 

268 @strawberry.type 

269 class Query: 

270 @trace_resolver("fetch_user") 

271 @strawberry.field 

272 async def user(self, info: Info, id: str) -> User: 

273 return await get_user(id) 

274 ``` 

275 """ 

276 

277 def decorator(func: Callable[P, T]) -> Callable[P, Any]: 

278 span_name = name or func.__name__ 

279 

280 @functools.wraps(func) 

281 async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: 

282 start_time = time.time() 

283 try: 

284 # Support coroutine functions but allow mypy to understand 

285 # that result could be awaitable or a direct value. 

286 result = func(*args, **kwargs) 

287 if isinstance(result, Awaitable): 

288 return cast("T", await result) 

289 return result 

290 finally: 

291 duration_ms = (time.time() - start_time) * 1000 

292 logger.debug("ResolverProtocol %s took %.2fms", span_name, duration_ms) 

293 

294 @functools.wraps(func) 

295 def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> T: 

296 start_time = time.time() 

297 try: 

298 return func(*args, **kwargs) 

299 finally: 

300 duration_ms = (time.time() - start_time) * 1000 

301 logger.debug("ResolverProtocol %s took %.2fms", span_name, duration_ms) 

302 

303 # Return appropriate wrapper 

304 import asyncio 

305 

306 if asyncio.iscoroutinefunction(func): 

307 return async_wrapper 

308 return sync_wrapper 

309 

310 return decorator 

311 

312 

313__all__ = [ 

314 "ExecutionTrace", 

315 "TraceSpan", 

316 "TracingExtension", 

317 "trace_resolver", 

318]