codekingpro/portable-devtools

1from types import SimpleNamespace
2from typing import TYPE_CHECKING, Mapping, Optional, Type, TypeVar
3
4import attr
5from aiosignal import Signal
6from multidict import CIMultiDict
7from yarl import URL
8
9from .client_reqrep import ClientResponse
10
11if TYPE_CHECKING:
12    from .client import ClientSession
13
14    _ParamT_contra = TypeVar("_ParamT_contra", contravariant=True)
15    _TracingSignal = Signal[ClientSession, SimpleNamespace, _ParamT_contra]
16
17
18__all__ = (
19    "TraceConfig",
20    "TraceRequestStartParams",
21    "TraceRequestEndParams",
22    "TraceRequestExceptionParams",
23    "TraceConnectionQueuedStartParams",
24    "TraceConnectionQueuedEndParams",
25    "TraceConnectionCreateStartParams",
26    "TraceConnectionCreateEndParams",
27    "TraceConnectionReuseconnParams",
28    "TraceDnsResolveHostStartParams",
29    "TraceDnsResolveHostEndParams",
30    "TraceDnsCacheHitParams",
31    "TraceDnsCacheMissParams",
32    "TraceRequestRedirectParams",
33    "TraceRequestChunkSentParams",
34    "TraceResponseChunkReceivedParams",
35    "TraceRequestHeadersSentParams",
36)
37
38
39class TraceConfig:
40    """First-class used to trace requests launched via ClientSession objects."""
41
42    def __init__(
43        self, trace_config_ctx_factory: Type[SimpleNamespace] = SimpleNamespace
44    ) -> None:
45        self._on_request_start: _TracingSignal[TraceRequestStartParams] = Signal(self)
46        self._on_request_chunk_sent: _TracingSignal[TraceRequestChunkSentParams] = (
47            Signal(self)
48        )
49        self._on_response_chunk_received: _TracingSignal[
50            TraceResponseChunkReceivedParams
51        ] = Signal(self)
52        self._on_request_end: _TracingSignal[TraceRequestEndParams] = Signal(self)
53        self._on_request_exception: _TracingSignal[TraceRequestExceptionParams] = (
54            Signal(self)
55        )
56        self._on_request_redirect: _TracingSignal[TraceRequestRedirectParams] = Signal(
57            self
58        )
59        self._on_connection_queued_start: _TracingSignal[
60            TraceConnectionQueuedStartParams
61        ] = Signal(self)
62        self._on_connection_queued_end: _TracingSignal[
63            TraceConnectionQueuedEndParams
64        ] = Signal(self)
65        self._on_connection_create_start: _TracingSignal[
66            TraceConnectionCreateStartParams
67        ] = Signal(self)
68        self._on_connection_create_end: _TracingSignal[
69            TraceConnectionCreateEndParams
70        ] = Signal(self)
71        self._on_connection_reuseconn: _TracingSignal[
72            TraceConnectionReuseconnParams
73        ] = Signal(self)
74        self._on_dns_resolvehost_start: _TracingSignal[
75            TraceDnsResolveHostStartParams
76        ] = Signal(self)
77        self._on_dns_resolvehost_end: _TracingSignal[TraceDnsResolveHostEndParams] = (
78            Signal(self)
79        )
80        self._on_dns_cache_hit: _TracingSignal[TraceDnsCacheHitParams] = Signal(self)
81        self._on_dns_cache_miss: _TracingSignal[TraceDnsCacheMissParams] = Signal(self)
82        self._on_request_headers_sent: _TracingSignal[TraceRequestHeadersSentParams] = (
83            Signal(self)
84        )
85
86        self._trace_config_ctx_factory = trace_config_ctx_factory
87
88    def trace_config_ctx(
89        self, trace_request_ctx: Optional[Mapping[str, str]] = None
90    ) -> SimpleNamespace:
91        """Return a new trace_config_ctx instance"""
92        return self._trace_config_ctx_factory(trace_request_ctx=trace_request_ctx)
93
94    def freeze(self) -> None:
95        self._on_request_start.freeze()
96        self._on_request_chunk_sent.freeze()
97        self._on_response_chunk_received.freeze()
98        self._on_request_end.freeze()
99        self._on_request_exception.freeze()
100        self._on_request_redirect.freeze()
101        self._on_connection_queued_start.freeze()
102        self._on_connection_queued_end.freeze()
103        self._on_connection_create_start.freeze()
104        self._on_connection_create_end.freeze()
105        self._on_connection_reuseconn.freeze()
106        self._on_dns_resolvehost_start.freeze()
107        self._on_dns_resolvehost_end.freeze()
108        self._on_dns_cache_hit.freeze()
109        self._on_dns_cache_miss.freeze()
110        self._on_request_headers_sent.freeze()
111
112    @property
113    def on_request_start(self) -> "_TracingSignal[TraceRequestStartParams]":
114        return self._on_request_start
115
116    @property
117    def on_request_chunk_sent(
118        self,
119    ) -> "_TracingSignal[TraceRequestChunkSentParams]":
120        return self._on_request_chunk_sent
121
122    @property
123    def on_response_chunk_received(
124        self,
125    ) -> "_TracingSignal[TraceResponseChunkReceivedParams]":
126        return self._on_response_chunk_received
127
128    @property
129    def on_request_end(self) -> "_TracingSignal[TraceRequestEndParams]":
130        return self._on_request_end
131
132    @property
133    def on_request_exception(
134        self,
135    ) -> "_TracingSignal[TraceRequestExceptionParams]":
136        return self._on_request_exception
137
138    @property
139    def on_request_redirect(
140        self,
141    ) -> "_TracingSignal[TraceRequestRedirectParams]":
142        return self._on_request_redirect
143
144    @property
145    def on_connection_queued_start(
146        self,
147    ) -> "_TracingSignal[TraceConnectionQueuedStartParams]":
148        return self._on_connection_queued_start
149
150    @property
151    def on_connection_queued_end(
152        self,
153    ) -> "_TracingSignal[TraceConnectionQueuedEndParams]":
154        return self._on_connection_queued_end
155
156    @property
157    def on_connection_create_start(
158        self,
159    ) -> "_TracingSignal[TraceConnectionCreateStartParams]":
160        return self._on_connection_create_start
161
162    @property
163    def on_connection_create_end(
164        self,
165    ) -> "_TracingSignal[TraceConnectionCreateEndParams]":
166        return self._on_connection_create_end
167
168    @property
169    def on_connection_reuseconn(
170        self,
171    ) -> "_TracingSignal[TraceConnectionReuseconnParams]":
172        return self._on_connection_reuseconn
173
174    @property
175    def on_dns_resolvehost_start(
176        self,
177    ) -> "_TracingSignal[TraceDnsResolveHostStartParams]":
178        return self._on_dns_resolvehost_start
179
180    @property
181    def on_dns_resolvehost_end(
182        self,
183    ) -> "_TracingSignal[TraceDnsResolveHostEndParams]":
184        return self._on_dns_resolvehost_end
185
186    @property
187    def on_dns_cache_hit(self) -> "_TracingSignal[TraceDnsCacheHitParams]":
188        return self._on_dns_cache_hit
189
190    @property
191    def on_dns_cache_miss(self) -> "_TracingSignal[TraceDnsCacheMissParams]":
192        return self._on_dns_cache_miss
193
194    @property
195    def on_request_headers_sent(
196        self,
197    ) -> "_TracingSignal[TraceRequestHeadersSentParams]":
198        return self._on_request_headers_sent
199
200
201@attr.s(auto_attribs=True, frozen=True, slots=True)
202class TraceRequestStartParams:
203    """Parameters sent by the `on_request_start` signal"""
204
205    method: str
206    url: URL
207    headers: "CIMultiDict[str]"
208
209
210@attr.s(auto_attribs=True, frozen=True, slots=True)
211class TraceRequestChunkSentParams:
212    """Parameters sent by the `on_request_chunk_sent` signal"""
213
214    method: str
215    url: URL
216    chunk: bytes
217
218
219@attr.s(auto_attribs=True, frozen=True, slots=True)
220class TraceResponseChunkReceivedParams:
221    """Parameters sent by the `on_response_chunk_received` signal"""
222
223    method: str
224    url: URL
225    chunk: bytes
226
227
228@attr.s(auto_attribs=True, frozen=True, slots=True)
229class TraceRequestEndParams:
230    """Parameters sent by the `on_request_end` signal"""
231
232    method: str
233    url: URL
234    headers: "CIMultiDict[str]"
235    response: ClientResponse
236
237
238@attr.s(auto_attribs=True, frozen=True, slots=True)
239class TraceRequestExceptionParams:
240    """Parameters sent by the `on_request_exception` signal"""
241
242    method: str
243    url: URL
244    headers: "CIMultiDict[str]"
245    exception: BaseException
246
247
248@attr.s(auto_attribs=True, frozen=True, slots=True)
249class TraceRequestRedirectParams:
250    """Parameters sent by the `on_request_redirect` signal"""
251
252    method: str
253    url: URL
254    headers: "CIMultiDict[str]"
255    response: ClientResponse
256
257
258@attr.s(auto_attribs=True, frozen=True, slots=True)
259class TraceConnectionQueuedStartParams:
260    """Parameters sent by the `on_connection_queued_start` signal"""
261
262
263@attr.s(auto_attribs=True, frozen=True, slots=True)
264class TraceConnectionQueuedEndParams:
265    """Parameters sent by the `on_connection_queued_end` signal"""
266
267
268@attr.s(auto_attribs=True, frozen=True, slots=True)
269class TraceConnectionCreateStartParams:
270    """Parameters sent by the `on_connection_create_start` signal"""
271
272
273@attr.s(auto_attribs=True, frozen=True, slots=True)
274class TraceConnectionCreateEndParams:
275    """Parameters sent by the `on_connection_create_end` signal"""
276
277
278@attr.s(auto_attribs=True, frozen=True, slots=True)
279class TraceConnectionReuseconnParams:
280    """Parameters sent by the `on_connection_reuseconn` signal"""
281
282
283@attr.s(auto_attribs=True, frozen=True, slots=True)
284class TraceDnsResolveHostStartParams:
285    """Parameters sent by the `on_dns_resolvehost_start` signal"""
286
287    host: str
288
289
290@attr.s(auto_attribs=True, frozen=True, slots=True)
291class TraceDnsResolveHostEndParams:
292    """Parameters sent by the `on_dns_resolvehost_end` signal"""
293
294    host: str
295
296
297@attr.s(auto_attribs=True, frozen=True, slots=True)
298class TraceDnsCacheHitParams:
299    """Parameters sent by the `on_dns_cache_hit` signal"""
300
301    host: str
302
303
304@attr.s(auto_attribs=True, frozen=True, slots=True)
305class TraceDnsCacheMissParams:
306    """Parameters sent by the `on_dns_cache_miss` signal"""
307
308    host: str
309
310
311@attr.s(auto_attribs=True, frozen=True, slots=True)
312class TraceRequestHeadersSentParams:
313    """Parameters sent by the `on_request_headers_sent` signal"""
314
315    method: str
316    url: URL
317    headers: "CIMultiDict[str]"
318
319
320class Trace:
321    """Internal dependency holder class.
322
323    Used to keep together the main dependencies used
324    at the moment of send a signal.
325    """
326
327    def __init__(
328        self,
329        session: "ClientSession",
330        trace_config: TraceConfig,
331        trace_config_ctx: SimpleNamespace,
332    ) -> None:
333        self._trace_config = trace_config
334        self._trace_config_ctx = trace_config_ctx
335        self._session = session
336
337    async def send_request_start(
338        self, method: str, url: URL, headers: "CIMultiDict[str]"
339    ) -> None:
340        return await self._trace_config.on_request_start.send(
341            self._session,
342            self._trace_config_ctx,
343            TraceRequestStartParams(method, url, headers),
344        )
345
346    async def send_request_chunk_sent(
347        self, method: str, url: URL, chunk: bytes
348    ) -> None:
349        return await self._trace_config.on_request_chunk_sent.send(
350            self._session,
351            self._trace_config_ctx,
352            TraceRequestChunkSentParams(method, url, chunk),
353        )
354
355    async def send_response_chunk_received(
356        self, method: str, url: URL, chunk: bytes
357    ) -> None:
358        return await self._trace_config.on_response_chunk_received.send(
359            self._session,
360            self._trace_config_ctx,
361            TraceResponseChunkReceivedParams(method, url, chunk),
362        )
363
364    async def send_request_end(
365        self,
366        method: str,
367        url: URL,
368        headers: "CIMultiDict[str]",
369        response: ClientResponse,
370    ) -> None:
371        return await self._trace_config.on_request_end.send(
372            self._session,
373            self._trace_config_ctx,
374            TraceRequestEndParams(method, url, headers, response),
375        )
376
377    async def send_request_exception(
378        self,
379        method: str,
380        url: URL,
381        headers: "CIMultiDict[str]",
382        exception: BaseException,
383    ) -> None:
384        return await self._trace_config.on_request_exception.send(
385            self._session,
386            self._trace_config_ctx,
387            TraceRequestExceptionParams(method, url, headers, exception),
388        )
389
390    async def send_request_redirect(
391        self,
392        method: str,
393        url: URL,
394        headers: "CIMultiDict[str]",
395        response: ClientResponse,
396    ) -> None:
397        return await self._trace_config._on_request_redirect.send(
398            self._session,
399            self._trace_config_ctx,
400            TraceRequestRedirectParams(method, url, headers, response),
401        )
402
403    async def send_connection_queued_start(self) -> None:
404        return await self._trace_config.on_connection_queued_start.send(
405            self._session, self._trace_config_ctx, TraceConnectionQueuedStartParams()
406        )
407
408    async def send_connection_queued_end(self) -> None:
409        return await self._trace_config.on_connection_queued_end.send(
410            self._session, self._trace_config_ctx, TraceConnectionQueuedEndParams()
411        )
412
413    async def send_connection_create_start(self) -> None:
414        return await self._trace_config.on_connection_create_start.send(
415            self._session, self._trace_config_ctx, TraceConnectionCreateStartParams()
416        )
417
418    async def send_connection_create_end(self) -> None:
419        return await self._trace_config.on_connection_create_end.send(
420            self._session, self._trace_config_ctx, TraceConnectionCreateEndParams()
421        )
422
423    async def send_connection_reuseconn(self) -> None:
424        return await self._trace_config.on_connection_reuseconn.send(
425            self._session, self._trace_config_ctx, TraceConnectionReuseconnParams()
426        )
427
428    async def send_dns_resolvehost_start(self, host: str) -> None:
429        return await self._trace_config.on_dns_resolvehost_start.send(
430            self._session, self._trace_config_ctx, TraceDnsResolveHostStartParams(host)
431        )
432
433    async def send_dns_resolvehost_end(self, host: str) -> None:
434        return await self._trace_config.on_dns_resolvehost_end.send(
435            self._session, self._trace_config_ctx, TraceDnsResolveHostEndParams(host)
436        )
437
438    async def send_dns_cache_hit(self, host: str) -> None:
439        return await self._trace_config.on_dns_cache_hit.send(
440            self._session, self._trace_config_ctx, TraceDnsCacheHitParams(host)
441        )
442
443    async def send_dns_cache_miss(self, host: str) -> None:
444        return await self._trace_config.on_dns_cache_miss.send(
445            self._session, self._trace_config_ctx, TraceDnsCacheMissParams(host)
446        )
447
448    async def send_request_headers(
449        self, method: str, url: URL, headers: "CIMultiDict[str]"
450    ) -> None:
451        return await self._trace_config._on_request_headers_sent.send(
452            self._session,
453            self._trace_config_ctx,
454            TraceRequestHeadersSentParams(method, url, headers),
455        )
456