def test_marshal_trace_context(self):
trace_id = "123456"
parent_id = "654321"
trace_fields = {"i": "like", "to": "trace"}
trace_context = marshal_trace_context(trace_id, parent_id, trace_fields)
trace_id_u, parent_id_u, trace_fields_u = unmarshal_trace_context(trace_context)
self.assertEqual(trace_id_u, trace_id, "unmarshaled trace id should match original")
self.assertEqual(parent_id_u, parent_id, "unmarshaled parent id should match original")
self.assertDictEqual(trace_fields_u, trace_fields, "unmarshaled trace fields should match original")
def _get_trace_data_from_message_attributes(attributes):
''' Look for the trace data from SNS/SQS Messsage Atrributes
'''
trace_id, parent_id, context = None, None, None
if isinstance(attributes, dict):
keymap = {k.lower(): k for k in attributes.keys()}
if 'x-honeycomb-trace' in keymap:
if 'Value' in attributes[keymap['x-honeycomb-trace']]:
# SNS
trace_id, parent_id, context = unmarshal_trace_context(
attributes[keymap['x-honeycomb-trace']]['Value']
)
elif 'stringValue' in attributes[keymap['x-honeycomb-trace']]:
# SQS
trace_id, parent_id, context = unmarshal_trace_context(
attributes[keymap['x-honeycomb-trace']]['stringValue']
)
return trace_id, parent_id, context
def _get_trace_context(request):
trace_context = request.META.get('HTTP_X_HONEYCOMB_TRACE')
beeline.internal.log("got trace context: %s", trace_context)
if trace_context:
try:
return unmarshal_trace_context(trace_context)
except Exception as e:
beeline.internal.log(
'error attempting to extract trace context: %s', str(e))
return None, None, None
def test_marshal_trace_context_empty_context(self):
trace_id = "123456"
parent_id = "654321"
trace_context = "{};trace_id={},parent_id={}".format(
1, trace_id, parent_id)
trace_id_u, parent_id_u, trace_fields_u = unmarshal_trace_context(
trace_context)
self.assertEqual(trace_id_u, trace_id,
"unmarshaled trace id should match original")
self.assertEqual(parent_id_u, parent_id,
"unmarshaled parent id should match original")
self.assertDictEqual(trace_fields_u, {},
"unmarshaled trace fields should match original")
def _get_trace_context(environ):
''' returns trace_id, parent_id, context '''
# http://werkzeug.pocoo.org/docs/0.14/wrappers/#base-wrappers
req = Request(environ, shallow=True)
trace_context = req.headers.get('x-honeycomb-trace')
beeline.internal.log("got trace context: %s", trace_context)
if trace_context:
try:
return unmarshal_trace_context(trace_context)
except Exception as e:
beeline.internal.log('error attempting to extract trace context: %s', beeline.internal.stringify_exception(e))
return None, None, None
def _get_trace_data(event):
''' Extract trace/parent ids and context object that are threaded through
in various ways from other beelines'''
trace_id, parent_id, context = None, None, None
# If API gateway is triggering the Lambda, the event will have headers
# and we can look for our trace headers
if isinstance(event, dict):
if 'headers' in event:
if isinstance(event['headers'], dict):
# deal with possible case issues
keymap = {k.lower(): k for k in event['headers'].keys()}
if 'x-honeycomb-trace' in keymap:
trace_id, parent_id, context = unmarshal_trace_context(
event['headers'][keymap['x-honeycomb-trace']])
return trace_id, parent_id, context
def _get_trace_data(event):
''' Extract trace/parent ids and context object that are threaded through
in various ways from other beelines'''
trace_id, parent_id, context = None, None, None
# Look for trace headers in common places
if isinstance(event, dict):
# If API gateway is triggering the Lambda, the event will have headers
# and we can look for our trace headers
# https://docs.aws.amazon.com/lambda/latest/dg/with-on-demand-https.html
if 'headers' in event:
if isinstance(event['headers'], dict):
# deal with possible case issues
keymap = {k.lower(): k for k in event['headers'].keys()}
if 'x-honeycomb-trace' in keymap:
trace_id, parent_id, context = unmarshal_trace_context(
event['headers'][keymap['x-honeycomb-trace']]
)
# If a message source is triggering the Lambda, the event may have
# our trace data in the message attributes
elif 'Records' in event:
# Only process batches of exactly 1
# Higher batch sizes would have multiple messages thus
# generating multiple traces and requiring manual instrumentation
if len(event['Records']) == 1:
# If SNS is triggering the Lambda
# https://docs.aws.amazon.com/lambda/latest/dg/with-sns.html
if 'EventSource' in event['Records'][0]:
if event['Records'][0]['EventSource'] == 'aws:sns':
trace_id, parent_id, context = _get_trace_data_from_message_attributes(
event['Records'][0]['Sns']['MessageAttributes']
)
# If SQS is triggering the Lambda
# https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html
elif 'eventSource' in event['Records'][0]:
if event['Records'][0]['eventSource'] == 'aws:sqs':
trace_id, parent_id, context = _get_trace_data_from_message_attributes(
event['Records'][0]['messageAttributes']
)
return trace_id, parent_id, context
def test_marshal_trace_context_dataset_included(self):
""" ensures unmarshalling still works if there's a dataset context field """
trace_id = "123456"
parent_id = "654321"
dataset_id = "foo"
trace_fields = {"i": "like", "to": "trace"}
fields_string = base64.b64encode(
json.dumps(trace_fields).encode()).decode()
trace_context = "{};trace_id={},parent_id={},dataset_id={},context={}".format(
1, trace_id, parent_id, dataset_id, fields_string)
trace_id_u, parent_id_u, trace_fields_u = unmarshal_trace_context(
trace_context)
self.assertEqual(trace_id_u, trace_id,
"unmarshaled trace id should match original")
self.assertEqual(parent_id_u, parent_id,
"unmarshaled parent id should match original")
self.assertDictEqual(trace_fields_u, trace_fields,
"unmarshaled trace fields should match original")
