Coverage for /pythoncovmergedfiles/medio/medio/src/aiohttp/aiohttp/tracing.py: 64%

1import dataclasses 

2from types import SimpleNamespace 

3from typing import TYPE_CHECKING, Awaitable, Optional, Type, TypeVar 

4 

5from aiosignal import Signal 

6from multidict import CIMultiDict 

7from yarl import URL 

8 

9from .client_reqrep import ClientResponse 

10 

11if TYPE_CHECKING: # pragma: no cover 

12 from typing_extensions import Protocol 

13 

14 from .client import ClientSession 

15 

16 _ParamT_contra = TypeVar("_ParamT_contra", contravariant=True) 

17 

18 class _SignalCallback(Protocol[_ParamT_contra]): 

19 def __call__( 

20 self, 

21 __client_session: ClientSession, 

22 __trace_config_ctx: SimpleNamespace, 

23 __params: _ParamT_contra, 

24 ) -> Awaitable[None]: 

25 ... 

26 

27 

28__all__ = ( 

29 "TraceConfig", 

30 "TraceRequestStartParams", 

31 "TraceRequestEndParams", 

32 "TraceRequestExceptionParams", 

33 "TraceConnectionQueuedStartParams", 

34 "TraceConnectionQueuedEndParams", 

35 "TraceConnectionCreateStartParams", 

36 "TraceConnectionCreateEndParams", 

37 "TraceConnectionReuseconnParams", 

38 "TraceDnsResolveHostStartParams", 

39 "TraceDnsResolveHostEndParams", 

40 "TraceDnsCacheHitParams", 

41 "TraceDnsCacheMissParams", 

42 "TraceRequestRedirectParams", 

43 "TraceRequestChunkSentParams", 

44 "TraceResponseChunkReceivedParams", 

45 "TraceRequestHeadersSentParams", 

46) 

47 

48 

49class TraceConfig: 

50 """First-class used to trace requests launched via ClientSession objects.""" 

51 

52 def __init__( 

53 self, trace_config_ctx_factory: Type[SimpleNamespace] = SimpleNamespace 

54 ) -> None: 

55 self._on_request_start: Signal[ 

56 _SignalCallback[TraceRequestStartParams] 

57 ] = Signal(self) 

58 self._on_request_chunk_sent: Signal[ 

59 _SignalCallback[TraceRequestChunkSentParams] 

60 ] = Signal(self) 

61 self._on_response_chunk_received: Signal[ 

62 _SignalCallback[TraceResponseChunkReceivedParams] 

63 ] = Signal(self) 

64 self._on_request_end: Signal[_SignalCallback[TraceRequestEndParams]] = Signal( 

65 self 

66 ) 

67 self._on_request_exception: Signal[ 

68 _SignalCallback[TraceRequestExceptionParams] 

69 ] = Signal(self) 

70 self._on_request_redirect: Signal[ 

71 _SignalCallback[TraceRequestRedirectParams] 

72 ] = Signal(self) 

73 self._on_connection_queued_start: Signal[ 

74 _SignalCallback[TraceConnectionQueuedStartParams] 

75 ] = Signal(self) 

76 self._on_connection_queued_end: Signal[ 

77 _SignalCallback[TraceConnectionQueuedEndParams] 

78 ] = Signal(self) 

79 self._on_connection_create_start: Signal[ 

80 _SignalCallback[TraceConnectionCreateStartParams] 

81 ] = Signal(self) 

82 self._on_connection_create_end: Signal[ 

83 _SignalCallback[TraceConnectionCreateEndParams] 

84 ] = Signal(self) 

85 self._on_connection_reuseconn: Signal[ 

86 _SignalCallback[TraceConnectionReuseconnParams] 

87 ] = Signal(self) 

88 self._on_dns_resolvehost_start: Signal[ 

89 _SignalCallback[TraceDnsResolveHostStartParams] 

90 ] = Signal(self) 

91 self._on_dns_resolvehost_end: Signal[ 

92 _SignalCallback[TraceDnsResolveHostEndParams] 

93 ] = Signal(self) 

94 self._on_dns_cache_hit: Signal[ 

95 _SignalCallback[TraceDnsCacheHitParams] 

96 ] = Signal(self) 

97 self._on_dns_cache_miss: Signal[ 

98 _SignalCallback[TraceDnsCacheMissParams] 

99 ] = Signal(self) 

100 self._on_request_headers_sent: Signal[ 

101 _SignalCallback[TraceRequestHeadersSentParams] 

102 ] = Signal(self) 

103 

104 self._trace_config_ctx_factory = trace_config_ctx_factory 

105 

106 def trace_config_ctx( 

107 self, trace_request_ctx: Optional[SimpleNamespace] = None 

108 ) -> SimpleNamespace: 

109 """Return a new trace_config_ctx instance""" 

110 return self._trace_config_ctx_factory(trace_request_ctx=trace_request_ctx) 

111 

112 def freeze(self) -> None: 

113 self._on_request_start.freeze() 

114 self._on_request_chunk_sent.freeze() 

115 self._on_response_chunk_received.freeze() 

116 self._on_request_end.freeze() 

117 self._on_request_exception.freeze() 

118 self._on_request_redirect.freeze() 

119 self._on_connection_queued_start.freeze() 

120 self._on_connection_queued_end.freeze() 

121 self._on_connection_create_start.freeze() 

122 self._on_connection_create_end.freeze() 

123 self._on_connection_reuseconn.freeze() 

124 self._on_dns_resolvehost_start.freeze() 

125 self._on_dns_resolvehost_end.freeze() 

126 self._on_dns_cache_hit.freeze() 

127 self._on_dns_cache_miss.freeze() 

128 self._on_request_headers_sent.freeze() 

129 

130 @property 

131 def on_request_start(self) -> "Signal[_SignalCallback[TraceRequestStartParams]]": 

132 return self._on_request_start 

133 

134 @property 

135 def on_request_chunk_sent( 

136 self, 

137 ) -> "Signal[_SignalCallback[TraceRequestChunkSentParams]]": 

138 return self._on_request_chunk_sent 

139 

140 @property 

141 def on_response_chunk_received( 

142 self, 

143 ) -> "Signal[_SignalCallback[TraceResponseChunkReceivedParams]]": 

144 return self._on_response_chunk_received 

145 

146 @property 

147 def on_request_end(self) -> "Signal[_SignalCallback[TraceRequestEndParams]]": 

148 return self._on_request_end 

149 

150 @property 

151 def on_request_exception( 

152 self, 

153 ) -> "Signal[_SignalCallback[TraceRequestExceptionParams]]": 

154 return self._on_request_exception 

155 

156 @property 

157 def on_request_redirect( 

158 self, 

159 ) -> "Signal[_SignalCallback[TraceRequestRedirectParams]]": 

160 return self._on_request_redirect 

161 

162 @property 

163 def on_connection_queued_start( 

164 self, 

165 ) -> "Signal[_SignalCallback[TraceConnectionQueuedStartParams]]": 

166 return self._on_connection_queued_start 

167 

168 @property 

169 def on_connection_queued_end( 

170 self, 

171 ) -> "Signal[_SignalCallback[TraceConnectionQueuedEndParams]]": 

172 return self._on_connection_queued_end 

173 

174 @property 

175 def on_connection_create_start( 

176 self, 

177 ) -> "Signal[_SignalCallback[TraceConnectionCreateStartParams]]": 

178 return self._on_connection_create_start 

179 

180 @property 

181 def on_connection_create_end( 

182 self, 

183 ) -> "Signal[_SignalCallback[TraceConnectionCreateEndParams]]": 

184 return self._on_connection_create_end 

185 

186 @property 

187 def on_connection_reuseconn( 

188 self, 

189 ) -> "Signal[_SignalCallback[TraceConnectionReuseconnParams]]": 

190 return self._on_connection_reuseconn 

191 

192 @property 

193 def on_dns_resolvehost_start( 

194 self, 

195 ) -> "Signal[_SignalCallback[TraceDnsResolveHostStartParams]]": 

196 return self._on_dns_resolvehost_start 

197 

198 @property 

199 def on_dns_resolvehost_end( 

200 self, 

201 ) -> "Signal[_SignalCallback[TraceDnsResolveHostEndParams]]": 

202 return self._on_dns_resolvehost_end 

203 

204 @property 

205 def on_dns_cache_hit(self) -> "Signal[_SignalCallback[TraceDnsCacheHitParams]]": 

206 return self._on_dns_cache_hit 

207 

208 @property 

209 def on_dns_cache_miss(self) -> "Signal[_SignalCallback[TraceDnsCacheMissParams]]": 

210 return self._on_dns_cache_miss 

211 

212 @property 

213 def on_request_headers_sent( 

214 self, 

215 ) -> "Signal[_SignalCallback[TraceRequestHeadersSentParams]]": 

216 return self._on_request_headers_sent 

217 

218 

219@dataclasses.dataclass(frozen=True) 

220class TraceRequestStartParams: 

221 """Parameters sent by the `on_request_start` signal""" 

222 

223 method: str 

224 url: URL 

225 headers: "CIMultiDict[str]" 

226 

227 

228@dataclasses.dataclass(frozen=True) 

229class TraceRequestChunkSentParams: 

230 """Parameters sent by the `on_request_chunk_sent` signal""" 

231 

232 method: str 

233 url: URL 

234 chunk: bytes 

235 

236 

237@dataclasses.dataclass(frozen=True) 

238class TraceResponseChunkReceivedParams: 

239 """Parameters sent by the `on_response_chunk_received` signal""" 

240 

241 method: str 

242 url: URL 

243 chunk: bytes 

244 

245 

246@dataclasses.dataclass(frozen=True) 

247class TraceRequestEndParams: 

248 """Parameters sent by the `on_request_end` signal""" 

249 

250 method: str 

251 url: URL 

252 headers: "CIMultiDict[str]" 

253 response: ClientResponse 

254 

255 

256@dataclasses.dataclass(frozen=True) 

257class TraceRequestExceptionParams: 

258 """Parameters sent by the `on_request_exception` signal""" 

259 

260 method: str 

261 url: URL 

262 headers: "CIMultiDict[str]" 

263 exception: BaseException 

264 

265 

266@dataclasses.dataclass(frozen=True) 

267class TraceRequestRedirectParams: 

268 """Parameters sent by the `on_request_redirect` signal""" 

269 

270 method: str 

271 url: URL 

272 headers: "CIMultiDict[str]" 

273 response: ClientResponse 

274 

275 

276@dataclasses.dataclass(frozen=True) 

277class TraceConnectionQueuedStartParams: 

278 """Parameters sent by the `on_connection_queued_start` signal""" 

279 

280 

281@dataclasses.dataclass(frozen=True) 

282class TraceConnectionQueuedEndParams: 

283 """Parameters sent by the `on_connection_queued_end` signal""" 

284 

285 

286@dataclasses.dataclass(frozen=True) 

287class TraceConnectionCreateStartParams: 

288 """Parameters sent by the `on_connection_create_start` signal""" 

289 

290 

291@dataclasses.dataclass(frozen=True) 

292class TraceConnectionCreateEndParams: 

293 """Parameters sent by the `on_connection_create_end` signal""" 

294 

295 

296@dataclasses.dataclass(frozen=True) 

297class TraceConnectionReuseconnParams: 

298 """Parameters sent by the `on_connection_reuseconn` signal""" 

299 

300 

301@dataclasses.dataclass(frozen=True) 

302class TraceDnsResolveHostStartParams: 

303 """Parameters sent by the `on_dns_resolvehost_start` signal""" 

304 

305 host: str 

306 

307 

308@dataclasses.dataclass(frozen=True) 

309class TraceDnsResolveHostEndParams: 

310 """Parameters sent by the `on_dns_resolvehost_end` signal""" 

311 

312 host: str 

313 

314 

315@dataclasses.dataclass(frozen=True) 

316class TraceDnsCacheHitParams: 

317 """Parameters sent by the `on_dns_cache_hit` signal""" 

318 

319 host: str 

320 

321 

322@dataclasses.dataclass(frozen=True) 

323class TraceDnsCacheMissParams: 

324 """Parameters sent by the `on_dns_cache_miss` signal""" 

325 

326 host: str 

327 

328 

329@dataclasses.dataclass(frozen=True) 

330class TraceRequestHeadersSentParams: 

331 """Parameters sent by the `on_request_headers_sent` signal""" 

332 

333 method: str 

334 url: URL 

335 headers: "CIMultiDict[str]" 

336 

337 

338class Trace: 

339 """Internal dependency holder class. 

340 

341 Used to keep together the main dependencies used 

342 at the moment of send a signal. 

343 """ 

344 

345 def __init__( 

346 self, 

347 session: "ClientSession", 

348 trace_config: TraceConfig, 

349 trace_config_ctx: SimpleNamespace, 

350 ) -> None: 

351 self._trace_config = trace_config 

352 self._trace_config_ctx = trace_config_ctx 

353 self._session = session 

354 

355 async def send_request_start( 

356 self, method: str, url: URL, headers: "CIMultiDict[str]" 

357 ) -> None: 

358 return await self._trace_config.on_request_start.send( 

359 self._session, 

360 self._trace_config_ctx, 

361 TraceRequestStartParams(method, url, headers), 

362 ) 

363 

364 async def send_request_chunk_sent( 

365 self, method: str, url: URL, chunk: bytes 

366 ) -> None: 

367 return await self._trace_config.on_request_chunk_sent.send( 

368 self._session, 

369 self._trace_config_ctx, 

370 TraceRequestChunkSentParams(method, url, chunk), 

371 ) 

372 

373 async def send_response_chunk_received( 

374 self, method: str, url: URL, chunk: bytes 

375 ) -> None: 

376 return await self._trace_config.on_response_chunk_received.send( 

377 self._session, 

378 self._trace_config_ctx, 

379 TraceResponseChunkReceivedParams(method, url, chunk), 

380 ) 

381 

382 async def send_request_end( 

383 self, 

384 method: str, 

385 url: URL, 

386 headers: "CIMultiDict[str]", 

387 response: ClientResponse, 

388 ) -> None: 

389 return await self._trace_config.on_request_end.send( 

390 self._session, 

391 self._trace_config_ctx, 

392 TraceRequestEndParams(method, url, headers, response), 

393 ) 

394 

395 async def send_request_exception( 

396 self, 

397 method: str, 

398 url: URL, 

399 headers: "CIMultiDict[str]", 

400 exception: BaseException, 

401 ) -> None: 

402 return await self._trace_config.on_request_exception.send( 

403 self._session, 

404 self._trace_config_ctx, 

405 TraceRequestExceptionParams(method, url, headers, exception), 

406 ) 

407 

408 async def send_request_redirect( 

409 self, 

410 method: str, 

411 url: URL, 

412 headers: "CIMultiDict[str]", 

413 response: ClientResponse, 

414 ) -> None: 

415 return await self._trace_config._on_request_redirect.send( 

416 self._session, 

417 self._trace_config_ctx, 

418 TraceRequestRedirectParams(method, url, headers, response), 

419 ) 

420 

421 async def send_connection_queued_start(self) -> None: 

422 return await self._trace_config.on_connection_queued_start.send( 

423 self._session, self._trace_config_ctx, TraceConnectionQueuedStartParams() 

424 ) 

425 

426 async def send_connection_queued_end(self) -> None: 

427 return await self._trace_config.on_connection_queued_end.send( 

428 self._session, self._trace_config_ctx, TraceConnectionQueuedEndParams() 

429 ) 

430 

431 async def send_connection_create_start(self) -> None: 

432 return await self._trace_config.on_connection_create_start.send( 

433 self._session, self._trace_config_ctx, TraceConnectionCreateStartParams() 

434 ) 

435 

436 async def send_connection_create_end(self) -> None: 

437 return await self._trace_config.on_connection_create_end.send( 

438 self._session, self._trace_config_ctx, TraceConnectionCreateEndParams() 

439 ) 

440 

441 async def send_connection_reuseconn(self) -> None: 

442 return await self._trace_config.on_connection_reuseconn.send( 

443 self._session, self._trace_config_ctx, TraceConnectionReuseconnParams() 

444 ) 

445 

446 async def send_dns_resolvehost_start(self, host: str) -> None: 

447 return await self._trace_config.on_dns_resolvehost_start.send( 

448 self._session, self._trace_config_ctx, TraceDnsResolveHostStartParams(host) 

449 ) 

450 

451 async def send_dns_resolvehost_end(self, host: str) -> None: 

452 return await self._trace_config.on_dns_resolvehost_end.send( 

453 self._session, self._trace_config_ctx, TraceDnsResolveHostEndParams(host) 

454 ) 

455 

456 async def send_dns_cache_hit(self, host: str) -> None: 

457 return await self._trace_config.on_dns_cache_hit.send( 

458 self._session, self._trace_config_ctx, TraceDnsCacheHitParams(host) 

459 ) 

460 

461 async def send_dns_cache_miss(self, host: str) -> None: 

462 return await self._trace_config.on_dns_cache_miss.send( 

463 self._session, self._trace_config_ctx, TraceDnsCacheMissParams(host) 

464 ) 

465 

466 async def send_request_headers( 

467 self, method: str, url: URL, headers: "CIMultiDict[str]" 

468 ) -> None: 

469 return await self._trace_config._on_request_headers_sent.send( 

470 self._session, 

471 self._trace_config_ctx, 

472 TraceRequestHeadersSentParams(method, url, headers), 

473 )