-
Notifications
You must be signed in to change notification settings - Fork 142
/
patch.py
56 lines (39 loc) · 1.46 KB
/
patch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import wrapt
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.models import http
from aws_xray_sdk.ext.util import inject_trace_header, strip_url, get_hostname
def patch():
wrapt.wrap_function_wrapper(
'requests',
'Session.request',
_xray_traced_requests
)
wrapt.wrap_function_wrapper(
'requests',
'Session.prepare_request',
_inject_header
)
def _xray_traced_requests(wrapped, instance, args, kwargs):
url = kwargs.get('url') or args[1]
return xray_recorder.record_subsegment(
wrapped, instance, args, kwargs,
name=get_hostname(url),
namespace='remote',
meta_processor=requests_processor,
)
def _inject_header(wrapped, instance, args, kwargs):
request = args[0]
headers = getattr(request, 'headers', {})
inject_trace_header(headers, xray_recorder.current_subsegment())
setattr(request, 'headers', headers)
return wrapped(*args, **kwargs)
def requests_processor(wrapped, instance, args, kwargs,
return_value, exception, subsegment, stack):
method = kwargs.get('method') or args[0]
url = kwargs.get('url') or args[1]
subsegment.put_http_meta(http.METHOD, method)
subsegment.put_http_meta(http.URL, strip_url(url))
if return_value is not None:
subsegment.put_http_meta(http.STATUS, return_value.status_code)
elif exception:
subsegment.add_exception(exception, stack)