Extract span tags from triggering event

datadog_lambda/tracing.py

@@ -89,7 +89,7 @@ def _context_obj_to_headers(obj):
     }
 
 
-def extract_dd_trace_context(event):
+def extract_dd_trace_context(event, tags):
     """
     Extract Datadog trace context from the Lambda `event` object.
 
@@ -114,6 +114,9 @@ def extract_dd_trace_context(event):
             "parent-id": parent_id,
             "sampling-priority": sampling_priority,
         }
+        # Add tags into the datadog-metadata subsegment
+        if tags:
+            metadata["tags"] = tags
         xray_recorder.begin_subsegment(XraySubsegment.NAME)
         subsegment = xray_recorder.current_subsegment()
 

datadog_lambda/trigger.py

@@ -0,0 +1,182 @@
+# 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": "dynamodb",
+    "aws:kinesis": "kinesis",
+    "aws:s3": "s3",
+    "aws:sns": "sns",
+    "aws:sqs": "sqs",
+}
+
+
+def get_first_record(event):
+    records = event.get("Records")
+    if records and len(records) > 0:
+        return records[0]
+
+
+def read_event_source(event):
+    event_source = event.get("eventSource") or event.get("EventSource")
+    return EVENT_SOURCES.get(event_source)
+
+
+def get_event_source(event):
+    """
+    Attempts to determine the source of the trigger event
+    """
+    event_source = read_event_source(event)
+
+    request_context = event.get("requestContext")
+    if request_context and request_context.get("stage"):
+        event_source = "api-gateway"
+
+    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 = read_event_source(event_record)
+        if event_record.get("cf"):
+            event_source = "cloudfront"
+
+    return event_source
+
+
+def read_event_source_arn(event):
+    event_source_arn = event.get("eventSourceARN") or event.get("eventSourceArn")
+    return event_source_arn
+
+
+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]
+
+    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:aws:cloudfront::{}:distribution/{}".format(
+            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:aws:apigateway:{}::/restapis/{}/stages/{}".format(
+            region, request_context["apiId"], request_context["stage"]
+        )
+
+    # 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:aws:logs:{}:{}:log-group:{}".format(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 = read_event_source_arn(event)
+
+    event_record = get_first_record(event)
+    if event_record:
+        event_source_arn = read_event_source_arn(event_record)
+
+    if event_source_arn is None:
+        event_source_arn = parse_event_source_arn(source, event, context)
+
+    return event_source_arn
+
+
+def get_apigateway_http_tags(event):
+    """
+    Extracts HTTP tags from the triggering API Gateway event into HTTP facet tags
+    """
+    http_tags = {}
+    request_context = event.get("requestContext")
+    if request_context:
+        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 get_trigger_tags(event, context):
+    """
+    Parses the trigger event object to get tags to be added to the span metadata
+    """
+    trigger_tags = {}
+    event_source = get_event_source(event)
+    if event_source:
+        trigger_tags["trigger.event_source"] = event_source
+
+        event_source_arn = get_event_source_arn(event_source, event, context)
+        if event_source_arn:
+            trigger_tags["trigger.event_source_arn"] = event_source_arn
+
+    if event_source == "api-gateway":
+        trigger_tags.update(get_apigateway_http_tags(event))
+
+    return trigger_tags
+
+
+def set_apigateway_status_code_tag(span, response):
+    """
+    If the Lambda was triggered by API Gateway add the returned status code
+    as a tag to the function execution span
+    """
+    is_apigateway_trigger = (
+        span and span.get_tag("trigger.event_source") == "api-gateway"
+    )
+    if is_apigateway_trigger and response and response.get("statusCode"):
+        span.set_tag("http.status_code", response.get("statusCode"))

datadog_lambda/wrapper.py

@@ -23,6 +23,7 @@
     set_dd_trace_py_root,
     create_function_execution_span,
 )
+from datadog_lambda.trigger import get_trigger_tags, set_apigateway_status_code_tag
 
 logger = logging.getLogger(__name__)
 
@@ -104,7 +105,9 @@ def __call__(self, event, context, **kwargs):
         """Executes when the wrapped function gets called"""
         self._before(event, context)
         try:
-            return self.func(event, context, **kwargs)
+            response = self.func(event, context, **kwargs)
+            set_apigateway_status_code_tag(self.span, response)
+            return response
         except Exception:
             submit_errors_metric(context)
             if self.span:
@@ -118,8 +121,10 @@ def _before(self, event, context):
 
             set_cold_start()
             submit_invocations_metric(context)
+            # Extract trigger tags
+            trigger_tags = get_trigger_tags(event, context)
             # Extract Datadog trace context from incoming requests
-            dd_context = extract_dd_trace_context(event)
+            dd_context = extract_dd_trace_context(event, trigger_tags)
 
             self.span = None
             if dd_tracing_enabled:
@@ -131,6 +136,9 @@ def _before(self, event, context):
                     dd_context,
                     self.merge_xray_traces,
                 )
+                # Add trigger tags
+                if self.span and trigger_tags:
+                    self.span.set_tags(trigger_tags)
             else:
                 set_correlation_ids()
 

tests/event_samples/api-gateway.json

@@ -0,0 +1,124 @@
+{
+    "body": "eyJ0ZXN0IjoiYm9keSJ9",
+    "resource": "/{proxy+}",
+    "path": "/path/to/resource",
+    "httpMethod": "POST",
+    "isBase64Encoded": true,
+    "queryStringParameters": {
+        "foo": "bar"
+    },
+    "multiValueQueryStringParameters": {
+        "foo": [
+            "bar"
+        ]
+    },
+    "pathParameters": {
+        "proxy": "/path/to/resource"
+    },
+    "stageVariables": {
+        "baz": "qux"
+    },
+    "headers": {
+        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
+        "Accept-Encoding": "gzip, deflate, sdch",
+        "Accept-Language": "en-US,en;q=0.8",
+        "Cache-Control": "max-age=0",
+        "CloudFront-Forwarded-Proto": "https",
+        "CloudFront-Is-Desktop-Viewer": "true",
+        "CloudFront-Is-Mobile-Viewer": "false",
+        "CloudFront-Is-SmartTV-Viewer": "false",
+        "CloudFront-Is-Tablet-Viewer": "false",
+        "CloudFront-Viewer-Country": "US",
+        "Host": "1234567890.execute-api.us-east-1.amazonaws.com",
+        "Upgrade-Insecure-Requests": "1",
+        "User-Agent": "Custom User Agent String",
+        "Via": "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)",
+        "X-Amz-Cf-Id": "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA==",
+        "X-Forwarded-For": "127.0.0.1, 127.0.0.2",
+        "X-Forwarded-Port": "443",
+        "X-Forwarded-Proto": "https"
+    },
+    "multiValueHeaders": {
+        "Accept": [
+            "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"
+        ],
+        "Accept-Encoding": [
+            "gzip, deflate, sdch"
+        ],
+        "Accept-Language": [
+            "en-US,en;q=0.8"
+        ],
+        "Cache-Control": [
+            "max-age=0"
+        ],
+        "CloudFront-Forwarded-Proto": [
+            "https"
+        ],
+        "CloudFront-Is-Desktop-Viewer": [
+            "true"
+        ],
+        "CloudFront-Is-Mobile-Viewer": [
+            "false"
+        ],
+        "CloudFront-Is-SmartTV-Viewer": [
+            "false"
+        ],
+        "CloudFront-Is-Tablet-Viewer": [
+            "false"
+        ],
+        "CloudFront-Viewer-Country": [
+            "US"
+        ],
+        "Host": [
+            "0123456789.execute-api.us-east-1.amazonaws.com"
+        ],
+        "Upgrade-Insecure-Requests": [
+            "1"
+        ],
+        "User-Agent": [
+            "Custom User Agent String"
+        ],
+        "Via": [
+            "1.1 08f323deadbeefa7af34d5feb414ce27.cloudfront.net (CloudFront)"
+        ],
+        "X-Amz-Cf-Id": [
+            "cDehVQoZnx43VYQb9j2-nvCh-9z396Uhbp027Y2JvkCPNLmGJHqlaA=="
+        ],
+        "X-Forwarded-For": [
+            "127.0.0.1, 127.0.0.2"
+        ],
+        "X-Forwarded-Port": [
+            "443"
+        ],
+        "X-Forwarded-Proto": [
+            "https"
+        ]
+    },
+    "requestContext": {
+        "accountId": "123456789012",
+        "resourceId": "123456",
+        "stage": "prod",
+        "requestId": "c6af9ac6-7b61-11e6-9a41-93e8deadbeef",
+        "requestTime": "09/Apr/2015:12:34:56 +0000",
+        "requestTimeEpoch": 1428582896000,
+        "identity": {
+            "cognitoIdentityPoolId": null,
+            "accountId": null,
+            "cognitoIdentityId": null,
+            "caller": null,
+            "accessKey": null,
+            "sourceIp": "127.0.0.1",
+            "cognitoAuthenticationType": null,
+            "cognitoAuthenticationProvider": null,
+            "userArn": null,
+            "userAgent": "Custom User Agent String",
+            "user": null
+        },
+        "domainName": "70ixmpl4fl.execute-api.us-east-2.amazonaws.com",
+        "path": "/prod/path/to/resource",
+        "resourcePath": "/{proxy+}",
+        "httpMethod": "POST",
+        "apiId": "1234567890",
+        "protocol": "HTTP/1.1"
+    }
+}