Extract span tags from triggering event (#101)

add function trigger span tags

Author: DylanLovesCoffee

datadog_lambda/constants.py

@@ -22,7 +22,8 @@ class TraceHeader(object):
 # X-Ray subsegment to save Datadog trace metadata
 class XraySubsegment(object):
  NAME = "datadog-metadata"
- KEY = "trace"
+ TRACE_KEY = "trace"
+ LAMBDA_FUNCTION_TAGS_KEY = "lambda_function_tags"
  NAMESPACE = "datadog"
 
 

datadog_lambda/tracing.py

@@ -62,7 +62,6 @@ def _get_xray_trace_context():
  "trace-id": _convert_xray_trace_id(xray_trace_entity.trace_id),
  "parent-id": _convert_xray_entity_id(xray_trace_entity.id),
  "sampling-priority": _convert_xray_sampling(xray_trace_entity.sampled),
- "source": TraceContextSource.XRAY,
  }
 
 
@@ -90,6 +89,22 @@ def _context_obj_to_headers(obj):
  }
 
 
+def create_dd_dummy_metadata_subsegment(
+ subsegment_metadata_value, subsegment_metadata_key
+):
+ """
+ Create a Datadog subsegment to pass the Datadog trace context or Lambda function
+ tags into its metadata field, so the X-Ray trace can be converted to a Datadog
+ trace in the Datadog backend with the correct context.
+ """
+ xray_recorder.begin_subsegment(XraySubsegment.NAME)
+ subsegment = xray_recorder.current_subsegment()
+ subsegment.put_metadata(
+ subsegment_metadata_key, subsegment_metadata_value, XraySubsegment.NAMESPACE
+ )
+ xray_recorder.end_subsegment()
+
+
 def extract_context_from_lambda_context(lambda_context):
  """
  Extract Datadog trace context from the `client_context` attr
@@ -156,12 +171,9 @@ def extract_dd_trace_context(event, lambda_context):
 
  Write the context to a global `dd_trace_context`, so the trace
  can be continued on the outgoing requests with the context injected.
-
- Save the context to an X-Ray subsegment's metadata field, so the X-Ray
- trace can be converted to a Datadog trace in the Datadog backend with
- the correct context.
  """
  global dd_trace_context
+ trace_context_source = None
 
  if "headers" in event:
  (
@@ -187,19 +199,16 @@ def extract_dd_trace_context(event, lambda_context):
  "parent-id": parent_id,
  "sampling-priority": sampling_priority,
  }
- xray_recorder.begin_subsegment(XraySubsegment.NAME)
- subsegment = xray_recorder.current_subsegment()
-
- subsegment.put_metadata(XraySubsegment.KEY, metadata, XraySubsegment.NAMESPACE)
  dd_trace_context = metadata.copy()
- dd_trace_context["source"] = TraceContextSource.EVENT
- xray_recorder.end_subsegment()
+ trace_context_source = TraceContextSource.EVENT
  else:
  # AWS Lambda runtime caches global variables between invocations,
  # reset to avoid using the context from the last invocation.
  dd_trace_context = _get_xray_trace_context()
+ if dd_trace_context:
+ trace_context_source = TraceContextSource.XRAY
  logger.debug("extracted dd trace context %s", dd_trace_context)
- return dd_trace_context
+ return dd_trace_context, trace_context_source
 
 
 def get_dd_trace_context():
@@ -300,15 +309,20 @@ def is_lambda_context():
  return type(xray_recorder.context) == LambdaContext
 
 
-def set_dd_trace_py_root(trace_context, merge_xray_traces):
- if trace_context["source"] == TraceContextSource.EVENT or merge_xray_traces:
+def set_dd_trace_py_root(trace_context_source, merge_xray_traces):
+ if trace_context_source == TraceContextSource.EVENT or merge_xray_traces:
  headers = get_dd_trace_context()
  span_context = propagator.extract(headers)
  tracer.context_provider.activate(span_context)
 
 
 def create_function_execution_span(
- context, function_name, is_cold_start, trace_context, merge_xray_traces
+ context,
+ function_name,
+ is_cold_start,
+ trace_context_source,
+ merge_xray_traces,
+ trigger_tags,
 ):
  tags = {}
  if context:
@@ -325,10 +339,9 @@ def create_function_execution_span(
  "datadog_lambda": datadog_lambda_version,
  "dd_trace": ddtrace_version,
  }
- source = trace_context["source"]
- if source == TraceContextSource.XRAY and merge_xray_traces:
- tags["_dd.parent_source"] = source
-
+ if trace_context_source == TraceContextSource.XRAY and merge_xray_traces:
+ tags["_dd.parent_source"] = trace_context_source
+ tags.update(trigger_tags)
  args = {
  "service": "aws.lambda",
  "resource": function_name,

datadog_lambda/trigger.py

@@ -0,0 +1,213 @@
+# 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 base64
+import gzip
+import json
+from io import BytesIO, BufferedReader
+
+
+EVENT_SOURCES = [
+ "aws:dynamodb",
+ "aws:kinesis",
+ "aws:s3",
+ "aws:sns",
+ "aws:sqs",
+]
+
+
+def get_aws_partition_by_region(region):
+ if region.startswith("us-gov-"):
+ return "aws-us-gov"
+ if region.startswith("cn-"):
+ return "aws-cn"
+ return "aws"
+
+
+def get_first_record(event):
+ records = event.get("Records")
+ if records and len(records) > 0:
+ return records[0]
+
+
+def parse_event_source(event):
+ """Determines the source of the trigger event
+
+ Possible Returns:
+ api-gateway | application-load-balancer | cloudwatch-logs |
+ cloudwatch-events | cloudfront | dynamodb | kinesis | s3 | sns | sqs
+ """
+ event_source = event.get("eventSource") or event.get("EventSource")
+
+ request_context = event.get("requestContext")
+ if request_context and request_context.get("stage"):
+ event_source = "api-gateway"
+
+ if request_context and request_context.get("elb"):
+ event_source = "application-load-balancer"
+
+ if event.get("awslogs"):
+ event_source = "cloudwatch-logs"
+
+ event_detail = event.get("detail")
+ cw_event_categories = event_detail and event_detail.get("EventCategories")
+ if event.get("source") == "aws.events" or cw_event_categories:
+ event_source = "cloudwatch-events"
+
+ event_record = get_first_record(event)
+ if event_record:
+ event_source = event_record.get("eventSource") or event_record.get(
+ "EventSource"
+ )
+ if event_record.get("cf"):
+ event_source = "cloudfront"
+
+ if event_source in EVENT_SOURCES:
+ event_source = event_source.replace("aws:", "")
+ return event_source
+
+
+def parse_event_source_arn(source, event, context):
+ """
+ Parses the trigger event for an available ARN. If an ARN field is not provided
+ in the event we stitch it together.
+ """
+ split_function_arn = context.invoked_function_arn.split(":")
+ region = split_function_arn[3]
+ account_id = split_function_arn[4]
+ aws_arn = get_aws_partition_by_region(region)
+
+ event_record = get_first_record(event)
+ # e.g. arn:aws:s3:::lambda-xyz123-abc890
+ if source == "s3":
+ return event_record.get("s3")["bucket"]["arn"]
+
+ # e.g. arn:aws:sns:us-east-1:123456789012:sns-lambda
+ if source == "sns":
+ return event_record.get("Sns")["TopicArn"]
+
+ # e.g. arn:aws:cloudfront::123456789012:distribution/ABC123XYZ
+ if source == "cloudfront":
+ distribution_id = event_record.get("cf")["config"]["distributionId"]
+ return "arn:{}:cloudfront::{}:distribution/{}".format(
+ aws_arn, account_id, distribution_id
+ )
+
+ # e.g. arn:aws:apigateway:us-east-1::/restapis/xyz123/stages/default
+ if source == "api-gateway":
+ request_context = event.get("requestContext")
+ return "arn:{}:apigateway:{}::/restapis/{}/stages/{}".format(
+ aws_arn, region, request_context["apiId"], request_context["stage"]
+ )
+
+ # e.g. arn:aws:elasticloadbalancing:us-east-1:123456789012:targetgroup/lambda-xyz/123
+ if source == "application-load-balancer":
+ request_context = event.get("requestContext")
+ return request_context.get("elb")["targetGroupArn"]
+
+ # e.g. arn:aws:logs:us-west-1:123456789012:log-group:/my-log-group-xyz
+ if source == "cloudwatch-logs":
+ with gzip.GzipFile(
+ fileobj=BytesIO(base64.b64decode(event["awslogs"]["data"]))
+ ) as decompress_stream:
+ data = b"".join(BufferedReader(decompress_stream))
+ logs = json.loads(data)
+ log_group = logs.get("logGroup", "cloudwatch")
+ return "arn:{}:logs:{}:{}:log-group:{}".format(
+ aws_arn, region, account_id, log_group
+ )
+
+ # e.g. arn:aws:events:us-east-1:123456789012:rule/my-schedule
+ if source == "cloudwatch-events" and event.get("resources"):
+ return event.get("resources")[0]
+
+
+def get_event_source_arn(source, event, context):
+ event_source_arn = event.get("eventSourceARN") or event.get("eventSourceArn")
+
+ event_record = get_first_record(event)
+ if event_record:
+ event_source_arn = event_record.get("eventSourceARN") or event_record.get(
+ "eventSourceArn"
+ )
+
+ if event_source_arn is None:
+ event_source_arn = parse_event_source_arn(source, event, context)
+
+ return event_source_arn
+
+
+def extract_http_tags(event):
+ """
+ Extracts HTTP facet tags from the triggering event
+ """
+ http_tags = {}
+ request_context = event.get("requestContext")
+ path = event.get("path")
+ method = event.get("httpMethod")
+ if request_context and request_context.get("stage"):
+ if request_context.get("domainName"):
+ http_tags["http.url"] = request_context["domainName"]
+
+ path = request_context.get("path")
+ method = request_context.get("httpMethod")
+ # Version 2.0 HTTP API Gateway
+ apigateway_v2_http = request_context.get("http")
+ if event.get("version") == "2.0" and apigateway_v2_http:
+ path = apigateway_v2_http.get("path")
+ method = apigateway_v2_http.get("method")
+
+ if path:
+ http_tags["http.url_details.path"] = path
+ if method:
+ http_tags["http.method"] = method
+
+ headers = event.get("headers")
+ if headers and headers.get("Referer"):
+ http_tags["http.referer"] = headers["Referer"]
+
+ return http_tags
+
+
+def extract_trigger_tags(event, context):
+ """
+ Parses the trigger event object to get tags to be added to the span metadata
+ """
+ trigger_tags = {}
+ event_source = parse_event_source(event)
+ if event_source:
+ trigger_tags["function_trigger.event_source"] = event_source
+
+ event_source_arn = get_event_source_arn(event_source, event, context)
+ if event_source_arn:
+ trigger_tags["function_trigger.event_source_arn"] = event_source_arn
+
+ if event_source in ["api-gateway", "application-load-balancer"]:
+ trigger_tags.update(extract_http_tags(event))
+
+ return trigger_tags
+
+
+def extract_http_status_code_tag(trigger_tags, response):
+ """
+ If the Lambda was triggered by API Gateway or ALB add the returned status code
+ as a tag to the function execution span.
+ """
+ is_http_trigger = trigger_tags and (
+ trigger_tags.get("function_trigger.event_source") == "api-gateway"
+ or trigger_tags.get("function_trigger.event_source")
+ == "application-load-balancer"
+ )
+ if not is_http_trigger:
+ return
+
+ status_code = "200"
+ if response is None:
+ # Return a 502 status if no response is found
+ status_code = "502"
+ elif response.get("statusCode"):
+ status_code = response.get("statusCode")
+
+ return status_code