Capture common HTTP attributes from API Gateway events in Lambda instrumentor

finlaysawyer · finlaysawyer · commit 576b96961acc · 2022-08-17T22:33:28.000+01:00
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py
@@ -64,11 +64,11 @@ def custom_event_context_extractor(lambda_event):
         event_context_extractor=custom_event_context_extractor
     )
 """
-
 import logging
 import os
 from importlib import import_module
 from typing import Any, Callable, Collection
+from urllib.parse import urlencode
 
 from wrapt import wrap_function_wrapper
 
@@ -85,6 +85,7 @@ def custom_event_context_extractor(lambda_event):
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import (
+    Span,
     SpanKind,
     TracerProvider,
     get_tracer,
@@ -171,6 +172,86 @@ def _determine_parent_context(
     return parent_context
 
 
+def _set_api_gateway_v1_proxy_attributes(
+    lambda_event: Any, span: Span
+) -> Span:
+    """Sets HTTP attributes for REST APIs and v1 HTTP APIs
+
+    More info:
+    https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format
+    """
+    span.set_attribute(
+        SpanAttributes.HTTP_METHOD, lambda_event.get("httpMethod")
+    )
+    span.set_attribute(SpanAttributes.HTTP_ROUTE, lambda_event.get("resource"))
+
+    if lambda_event.get("headers"):
+        span.set_attribute(
+            SpanAttributes.HTTP_USER_AGENT,
+            lambda_event["headers"].get("User-Agent"),
+        )
+        span.set_attribute(
+            SpanAttributes.HTTP_SCHEME,
+            lambda_event["headers"].get("X-Forwarded-Proto"),
+        )
+        span.set_attribute(
+            SpanAttributes.NET_HOST_NAME, lambda_event["headers"].get("Host")
+        )
+
+    if lambda_event.get("queryStringParameters"):
+        span.set_attribute(
+            SpanAttributes.HTTP_TARGET,
+            f"{lambda_event.get('resource')}?{urlencode(lambda_event.get('queryStringParameters'))}",
+        )
+    else:
+        span.set_attribute(
+            SpanAttributes.HTTP_TARGET, lambda_event.get("resource")
+        )
+
+    return span
+
+
+def _set_api_gateway_v2_proxy_attributes(
+    lambda_event: Any, span: Span
+) -> Span:
+    """Sets HTTP attributes for v2 HTTP APIs
+
+    More info:
+    https://docs.aws.amazon.com/apigateway/latest/developerguide/http-api-develop-integrations-lambda.html
+    """
+    span.set_attribute(
+        SpanAttributes.NET_HOST_NAME,
+        lambda_event["requestContext"].get("domainName"),
+    )
+
+    if lambda_event["requestContext"].get("http"):
+        span.set_attribute(
+            SpanAttributes.HTTP_METHOD,
+            lambda_event["requestContext"]["http"].get("method"),
+        )
+        span.set_attribute(
+            SpanAttributes.HTTP_USER_AGENT,
+            lambda_event["requestContext"]["http"].get("userAgent"),
+        )
+        span.set_attribute(
+            SpanAttributes.HTTP_ROUTE,
+            lambda_event["requestContext"]["http"].get("path"),
+        )
+
+        if lambda_event.get("rawQueryString"):
+            span.set_attribute(
+                SpanAttributes.HTTP_TARGET,
+                f"{lambda_event['requestContext']['http'].get('path')}?{lambda_event.get('rawQueryString')}",
+            )
+        else:
+            span.set_attribute(
+                SpanAttributes.HTTP_TARGET,
+                lambda_event["requestContext"]["http"].get("path"),
+            )
+
+    return span
+
+
 def _instrument(
     wrapped_module_name,
     wrapped_function_name,
@@ -233,6 +314,23 @@ def _instrumented_lambda_handler_call(
 
             result = call_wrapped(*args, **kwargs)
 
+            # If the request came from an API Gateway, extract http attributes from the event
+            # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#api-gateway
+            # https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/http.md#http-server-semantic-conventions
+            if lambda_event and lambda_event.get("requestContext"):
+                span.set_attribute(SpanAttributes.FAAS_TRIGGER, "http")
+
+                if lambda_event.get("version") == "2.0":
+                    _set_api_gateway_v2_proxy_attributes(lambda_event, span)
+                else:
+                    _set_api_gateway_v1_proxy_attributes(lambda_event, span)
+
+                if isinstance(result, dict) and result.get("statusCode"):
+                    span.set_attribute(
+                        SpanAttributes.HTTP_STATUS_CODE,
+                        result.get("statusCode"),
+                    )
+
         _tracer_provider = tracer_provider or get_tracer_provider()
         try:
             # NOTE: `force_flush` before function quit in case of Lambda freeze.
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/api_gateway_http_api_event.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/api_gateway_http_api_event.py
@@ -0,0 +1,54 @@
+# Generated via `sam local generate-event apigateway http-api-proxy`
+
+MOCK_LAMBDA_API_GATEWAY_HTTP_API_EVENT = {
+    "version": "2.0",
+    "routeKey": "$default",
+    "rawPath": "/path/to/resource",
+    "rawQueryString": "parameter1=value1&parameter1=value2&parameter2=value",
+    "cookies": ["cookie1", "cookie2"],
+    "headers": {"header1": "value1", "Header2": "value1,value2"},
+    "queryStringParameters": {
+        "parameter1": "value1,value2",
+        "parameter2": "value",
+    },
+    "requestContext": {
+        "accountId": "123456789012",
+        "apiId": "api-id",
+        "authentication": {
+            "clientCert": {
+                "clientCertPem": "CERT_CONTENT",
+                "subjectDN": "www.example.com",
+                "issuerDN": "Example issuer",
+                "serialNumber": "a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1:a1",
+                "validity": {
+                    "notBefore": "May 28 12:30:02 2019 GMT",
+                    "notAfter": "Aug  5 09:36:04 2021 GMT",
+                },
+            }
+        },
+        "authorizer": {
+            "jwt": {
+                "claims": {"claim1": "value1", "claim2": "value2"},
+                "scopes": ["scope1", "scope2"],
+            }
+        },
+        "domainName": "id.execute-api.us-east-1.amazonaws.com",
+        "domainPrefix": "id",
+        "http": {
+            "method": "POST",
+            "path": "/path/to/resource",
+            "protocol": "HTTP/1.1",
+            "sourceIp": "192.168.0.1/32",
+            "userAgent": "agent",
+        },
+        "requestId": "id",
+        "routeKey": "$default",
+        "stage": "$default",
+        "time": "12/Mar/2020:19:03:58 +0000",
+        "timeEpoch": 1583348638390,
+    },
+    "body": "eyJ0ZXN0IjoiYm9keSJ9",
+    "pathParameters": {"parameter1": "value1"},
+    "isBase64Encoded": True,
+    "stageVariables": {"stageVariable1": "value1", "stageVariable2": "value2"},
+}
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/api_gateway_proxy_event.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/api_gateway_proxy_event.py
@@ -0,0 +1,85 @@
+# Generated via `sam local generate-event apigateway aws-proxy`
+
+MOCK_LAMBDA_API_GATEWAY_PROXY_EVENT = {
+    "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": None,
+            "accountId": None,
+            "cognitoIdentityId": None,
+            "caller": None,
+            "accessKey": None,
+            "sourceIp": "127.0.0.1",
+            "cognitoAuthenticationType": None,
+            "cognitoAuthenticationProvider": None,
+            "userArn": None,
+            "userAgent": "Custom User Agent String",
+            "user": None,
+        },
+        "path": "/prod/path/to/resource",
+        "resourcePath": "/{proxy+}",
+        "httpMethod": "POST",
+        "apiId": "1234567890",
+        "protocol": "HTTP/1.1",
+    },
+}
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/lambda_function.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/mocks/lambda_function.py
@@ -15,3 +15,7 @@
 
 def handler(event, context):
     return "200 ok"
+
+
+def rest_api_handler(event, context):
+    return {"statusCode": 200, "body": "200 ok"}
diff --git a/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py b/instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py
@@ -15,6 +15,11 @@
 from importlib import import_module
 from unittest import mock
 
+from mocks.api_gateway_http_api_event import (
+    MOCK_LAMBDA_API_GATEWAY_HTTP_API_EVENT,
+)
+from mocks.api_gateway_proxy_event import MOCK_LAMBDA_API_GATEWAY_PROXY_EVENT
+
 from opentelemetry.environment_variables import OTEL_PROPAGATORS
 from opentelemetry.instrumentation.aws_lambda import (
     _HANDLER,
@@ -300,3 +305,49 @@ def test_lambda_handles_multiple_consumers(self):
         assert spans
 
         test_env_patch.stop()
+
+    def test_api_gateway_proxy_event_sets_attributes(self):
+        handler_patch = mock.patch.dict(
+            "os.environ",
+            {_HANDLER: "mocks.lambda_function.rest_api_handler"},
+        )
+        handler_patch.start()
+
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda(MOCK_LAMBDA_API_GATEWAY_PROXY_EVENT)
+
+        span = self.memory_exporter.get_finished_spans()[0]
+
+        self.assertSpanHasAttributes(
+            span,
+            {
+                SpanAttributes.FAAS_TRIGGER: "http",
+                SpanAttributes.HTTP_METHOD: "POST",
+                SpanAttributes.HTTP_ROUTE: "/{proxy+}",
+                SpanAttributes.HTTP_TARGET: "/{proxy+}?foo=bar",
+                SpanAttributes.NET_HOST_NAME: "1234567890.execute-api.us-east-1.amazonaws.com",
+                SpanAttributes.HTTP_USER_AGENT: "Custom User Agent String",
+                SpanAttributes.HTTP_SCHEME: "https",
+                SpanAttributes.HTTP_STATUS_CODE: 200,
+            },
+        )
+
+    def test_api_gateway_http_api_proxy_event_sets_attributes(self):
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda(MOCK_LAMBDA_API_GATEWAY_HTTP_API_EVENT)
+
+        span = self.memory_exporter.get_finished_spans()[0]
+
+        self.assertSpanHasAttributes(
+            span,
+            {
+                SpanAttributes.FAAS_TRIGGER: "http",
+                SpanAttributes.HTTP_METHOD: "POST",
+                SpanAttributes.HTTP_ROUTE: "/path/to/resource",
+                SpanAttributes.HTTP_TARGET: "/path/to/resource?parameter1=value1&parameter1=value2&parameter2=value",
+                SpanAttributes.NET_HOST_NAME: "id.execute-api.us-east-1.amazonaws.com",
+                SpanAttributes.HTTP_USER_AGENT: "agent",
+            },
+        )
