|
68 | 68 |
|
69 | 69 | propagator = HTTPPropagator()
|
70 | 70 |
|
| 71 | +DD_TRACE_JAVA_TRACE_ID_PADDING = "00000000" |
| 72 | + |
71 | 73 |
|
72 | 74 | def _convert_xray_trace_id(xray_trace_id):
|
73 | 75 | """
|
@@ -248,28 +250,52 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
|
248 | 250 | first_record.get("Sns", {}).get("MessageAttributes", {}),
|
249 | 251 | )
|
250 | 252 | dd_payload = msg_attributes.get("_datadog", {})
|
251 |
| - # SQS uses dataType and binaryValue/stringValue |
252 |
| - # SNS uses Type and Value |
253 |
| - dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", "")) |
254 |
| - if dd_json_data_type == "Binary": |
255 |
| - dd_json_data = dd_payload.get( |
256 |
| - "binaryValue", |
257 |
| - dd_payload.get("Value", r"{}"), |
258 |
| - ) |
259 |
| - dd_json_data = base64.b64decode(dd_json_data) |
260 |
| - elif dd_json_data_type == "String": |
261 |
| - dd_json_data = dd_payload.get( |
262 |
| - "stringValue", |
263 |
| - dd_payload.get("Value", r"{}"), |
264 |
| - ) |
| 253 | + if dd_payload: |
| 254 | + # SQS uses dataType and binaryValue/stringValue |
| 255 | + # SNS uses Type and Value |
| 256 | + dd_json_data = None |
| 257 | + dd_json_data_type = dd_payload.get("Type", dd_payload.get("dataType", "")) |
| 258 | + if dd_json_data_type == "Binary": |
| 259 | + dd_json_data = dd_payload.get( |
| 260 | + "binaryValue", |
| 261 | + dd_payload.get("Value", r"{}"), |
| 262 | + ) |
| 263 | + dd_json_data = base64.b64decode(dd_json_data) |
| 264 | + elif dd_json_data_type == "String": |
| 265 | + dd_json_data = dd_payload.get( |
| 266 | + "stringValue", |
| 267 | + dd_payload.get("Value", r"{}"), |
| 268 | + ) |
| 269 | + else: |
| 270 | + logger.debug( |
| 271 | + "Datadog Lambda Python only supports extracting trace" |
| 272 | + "context from String or Binary SQS/SNS message attributes" |
| 273 | + ) |
| 274 | + |
| 275 | + if dd_json_data: |
| 276 | + dd_data = json.loads(dd_json_data) |
| 277 | + return propagator.extract(dd_data) |
265 | 278 | else:
|
266 |
| - logger.debug( |
267 |
| - "Datadog Lambda Python only supports extracting trace" |
268 |
| - "context from String or Binary SQS/SNS message attributes" |
269 |
| - ) |
270 |
| - return extract_context_from_lambda_context(lambda_context) |
271 |
| - dd_data = json.loads(dd_json_data) |
272 |
| - return propagator.extract(dd_data) |
| 279 | + # Handle case where trace context is injected into attributes.AWSTraceHeader |
| 280 | + # example: Root=1-654321ab-000000001234567890abcdef;Parent=0123456789abcdef;Sampled=1 |
| 281 | + x_ray_header = first_record.get("attributes", {}).get("AWSTraceHeader") |
| 282 | + if x_ray_header: |
| 283 | + x_ray_context = parse_xray_header(x_ray_header) |
| 284 | + trace_id_parts = x_ray_context.get("trace_id", "").split("-") |
| 285 | + if len(trace_id_parts) > 2 and trace_id_parts[2].startswith( |
| 286 | + DD_TRACE_JAVA_TRACE_ID_PADDING |
| 287 | + ): |
| 288 | + # If it starts with eight 0's padding, |
| 289 | + # then this AWSTraceHeader contains Datadog injected trace context |
| 290 | + logger.debug( |
| 291 | + "Found dd-trace injected trace context from AWSTraceHeader" |
| 292 | + ) |
| 293 | + return Context( |
| 294 | + trace_id=int(trace_id_parts[2][8:], 16), |
| 295 | + span_id=int(int(x_ray_context["parent_id"], 16)), |
| 296 | + sampling_priority=float(x_ray_context["sampled"]), |
| 297 | + ) |
| 298 | + return extract_context_from_lambda_context(lambda_context) |
273 | 299 | except Exception as e:
|
274 | 300 | logger.debug("The trace extractor returned with error %s", e)
|
275 | 301 | return extract_context_from_lambda_context(lambda_context)
|
|
0 commit comments