forked from DataDog/datadog-lambda-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathpatch.py
More file actions
95 lines (79 loc) · 2.43 KB
/
patch.py
File metadata and controls
95 lines (79 loc) · 2.43 KB
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
# Unless explicitly stated otherwise all files in this repository are licensed
# under the Apache License Version 2.0.
# This product includes software developed at Datadog (https://www.datadoghq.com/).
# Copyright 2019 Datadog, Inc.
import sys
import logging
from wrapt import wrap_function_wrapper as wrap
from datadog_lambda.tracing import get_dd_trace_context
logger = logging.getLogger(__name__)
if sys.version_info >= (3, 0, 0):
httplib_module = 'http.client'
else:
httplib_module = 'httplib'
_httplib_patched = False
_requests_patched = False
def patch_all():
"""
Patch the widely-used HTTP clients to automatically inject
Datadog trace context.
"""
_patch_httplib()
_patch_requests()
def _patch_httplib():
"""
Patch the Python built-in `httplib` (Python 2) or
`http.client` (Python 3) module.
"""
global _httplib_patched
if not _httplib_patched:
_httplib_patched = True
wrap(
httplib_module,
'HTTPConnection.request',
_wrap_httplib_request
)
logger.debug('Patched %s', httplib_module)
def _patch_requests():
"""
Patch the high-level HTTP client module `requests`
if it's installed.
"""
global _requests_patched
if not _requests_patched:
_requests_patched = True
try:
wrap(
'requests',
'Session.request',
_wrap_requests_request
)
logger.debug('Patched requests')
except Exception:
logger.debug('Failed to patch requests', exc_info=True)
def _wrap_requests_request(func, instance, args, kwargs):
"""
Wrap `requests.Session.request` to inject the Datadog trace headers
into the outgoing requests.
"""
context = get_dd_trace_context()
if 'headers' in kwargs:
kwargs['headers'].update(context)
elif len(args) >= 5:
args[4].update(context)
else:
kwargs['headers'] = context
return func(*args, **kwargs)
def _wrap_httplib_request(func, instance, args, kwargs):
"""
Wrap `httplib` (python2) or `http.client` (python3) to inject
the Datadog trace headers into the outgoing requests.
"""
context = get_dd_trace_context()
if 'headers' in kwargs:
kwargs['headers'].update(context)
elif len(args) >= 4:
args[3].update(context)
else:
kwargs['headers'] = context
return func(*args, **kwargs)