9
9
from aws_xray_sdk .core .lambda_launcher import LambdaContext
10
10
11
11
from ddtrace import patch , tracer
12
- from datadog_lambda .constants import (
13
- SamplingPriority ,
14
- TraceHeader ,
15
- XraySubsegment ,
16
- )
12
+ from datadog_lambda .constants import SamplingPriority , TraceHeader , XraySubsegment
17
13
18
14
logger = logging .getLogger (__name__ )
19
15
@@ -38,8 +34,11 @@ def _convert_xray_sampling(xray_sampled):
38
34
"""
39
35
Convert X-Ray sampled (True/False) to its Datadog counterpart.
40
36
"""
41
- return str (SamplingPriority .USER_KEEP ) if xray_sampled \
37
+ return (
38
+ str (SamplingPriority .USER_KEEP )
39
+ if xray_sampled
42
40
else str (SamplingPriority .USER_REJECT )
41
+ )
43
42
44
43
45
44
def extract_dd_trace_context (event ):
@@ -54,33 +53,31 @@ def extract_dd_trace_context(event):
54
53
the correct context.
55
54
"""
56
55
global dd_trace_context
57
- headers = event .get (' headers' , {})
56
+ headers = event .get (" headers" , {})
58
57
lowercase_headers = {k .lower (): v for k , v in headers .items ()}
59
58
60
59
trace_id = lowercase_headers .get (TraceHeader .TRACE_ID )
61
60
parent_id = lowercase_headers .get (TraceHeader .PARENT_ID )
62
61
sampling_priority = lowercase_headers .get (TraceHeader .SAMPLING_PRIORITY )
63
62
if trace_id and parent_id and sampling_priority :
64
- logger .debug (' Extracted Datadog trace context from headers' )
63
+ logger .debug (" Extracted Datadog trace context from headers" )
65
64
dd_trace_context = {
66
- ' trace-id' : trace_id ,
67
- ' parent-id' : parent_id ,
68
- ' sampling-priority' : sampling_priority ,
65
+ " trace-id" : trace_id ,
66
+ " parent-id" : parent_id ,
67
+ " sampling-priority" : sampling_priority ,
69
68
}
70
69
xray_recorder .begin_subsegment (XraySubsegment .NAME )
71
70
subsegment = xray_recorder .current_subsegment ()
72
71
subsegment .put_metadata (
73
- XraySubsegment .KEY ,
74
- dd_trace_context ,
75
- XraySubsegment .NAMESPACE
72
+ XraySubsegment .KEY , dd_trace_context , XraySubsegment .NAMESPACE
76
73
)
77
74
xray_recorder .end_subsegment ()
78
75
else :
79
76
# AWS Lambda runtime caches global variables between invocations,
80
77
# reset to avoid using the context from the last invocation.
81
78
dd_trace_context = {}
82
79
83
- logger .debug (' extracted dd trace context %s' , dd_trace_context )
80
+ logger .debug (" extracted dd trace context %s" , dd_trace_context )
84
81
85
82
86
83
def get_dd_trace_context ():
@@ -96,28 +93,24 @@ def get_dd_trace_context():
96
93
context to an outgoing request.
97
94
"""
98
95
if not is_lambda_context ():
99
- logger .debug (' get_dd_trace_context is only supported in LambdaContext' )
96
+ logger .debug (" get_dd_trace_context is only supported in LambdaContext" )
100
97
return {}
101
98
102
99
global dd_trace_context
103
100
xray_trace_entity = xray_recorder .get_trace_entity () # xray (sub)segment
104
101
if dd_trace_context :
105
102
return {
106
- TraceHeader .TRACE_ID :
107
- dd_trace_context ['trace-id' ],
108
- TraceHeader .PARENT_ID : _convert_xray_entity_id (
109
- xray_trace_entity .id ),
110
- TraceHeader .SAMPLING_PRIORITY :
111
- dd_trace_context ['sampling-priority' ],
103
+ TraceHeader .TRACE_ID : dd_trace_context ["trace-id" ],
104
+ TraceHeader .PARENT_ID : _convert_xray_entity_id (xray_trace_entity .id ),
105
+ TraceHeader .SAMPLING_PRIORITY : dd_trace_context ["sampling-priority" ],
112
106
}
113
107
else :
114
108
return {
115
- TraceHeader .TRACE_ID : _convert_xray_trace_id (
116
- xray_trace_entity .trace_id ),
117
- TraceHeader .PARENT_ID : _convert_xray_entity_id (
118
- xray_trace_entity .id ),
109
+ TraceHeader .TRACE_ID : _convert_xray_trace_id (xray_trace_entity .trace_id ),
110
+ TraceHeader .PARENT_ID : _convert_xray_entity_id (xray_trace_entity .id ),
119
111
TraceHeader .SAMPLING_PRIORITY : _convert_xray_sampling (
120
- xray_trace_entity .sampled ),
112
+ xray_trace_entity .sampled
113
+ ),
121
114
}
122
115
123
116
@@ -130,16 +123,16 @@ def set_correlation_ids():
130
123
TODO: Remove me when Datadog tracer is natively supported in Lambda.
131
124
"""
132
125
if not is_lambda_context ():
133
- logger .debug (' set_correlation_ids is only supported in LambdaContext' )
126
+ logger .debug (" set_correlation_ids is only supported in LambdaContext" )
134
127
return
135
128
136
129
context = get_dd_trace_context ()
137
130
138
- span = tracer .trace (' dummy.span' )
131
+ span = tracer .trace (" dummy.span" )
139
132
span .trace_id = context [TraceHeader .TRACE_ID ]
140
133
span .span_id = context [TraceHeader .PARENT_ID ]
141
134
142
- logger .debug (' correlation ids set' )
135
+ logger .debug (" correlation ids set" )
143
136
144
137
145
138
def inject_correlation_ids ():
@@ -153,17 +146,19 @@ def inject_correlation_ids():
153
146
# Override the log format of the AWS provided LambdaLoggerHandler
154
147
root_logger = logging .getLogger ()
155
148
for handler in root_logger .handlers :
156
- if handler .__class__ .__name__ == 'LambdaLoggerHandler' :
157
- handler .setFormatter (logging .Formatter (
158
- '[%(levelname)s]\t %(asctime)s.%(msecs)dZ\t %(aws_request_id)s\t '
159
- '[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s]\t %(message)s\n ' ,
160
- '%Y-%m-%dT%H:%M:%S'
161
- ))
149
+ if handler .__class__ .__name__ == "LambdaLoggerHandler" :
150
+ handler .setFormatter (
151
+ logging .Formatter (
152
+ "[%(levelname)s]\t %(asctime)s.%(msecs)dZ\t %(aws_request_id)s\t "
153
+ "[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s]\t %(message)s\n " ,
154
+ "%Y-%m-%dT%H:%M:%S" ,
155
+ )
156
+ )
162
157
163
158
# Patch `logging.Logger.makeRecord` to actually inject correlation ids
164
159
patch (logging = True )
165
160
166
- logger .debug (' logs injection configured' )
161
+ logger .debug (" logs injection configured" )
167
162
168
163
169
164
def is_lambda_context ():
0 commit comments