Instrument aws elb serverless requests (#1605)

* add elb events processing. fix full url forming for api gw events

* add test for query parameters for api gateway v1

* fmt

* Small fix + CHANGELOG

* Fix origin.account_id vs origin.account.id for SNS

Co-authored-by: Anton Matviyenko <[email protected]>
Co-authored-by: Colton Myers <[email protected]>

Author: 3 people

CHANGELOG.asciidoc

@@ -29,6 +29,18 @@ endif::[]
 //===== Bug fixes
 //
 
+=== Unreleased
+
+// Unreleased changes go here
+// When the next release happens, nest these changes under the "Python Agent version 6.x" heading
+[float]
+===== Features
+
+* Added lambda support for ELB triggers {pull}#1605[#1605]
+
+//[float]
+//===== Bug fixes
+
 
 
 [[release-notes-6.x]]

elasticapm/contrib/serverless/aws.py

@@ -36,6 +36,7 @@
 import platform
 import time
 from typing import Optional
+from urllib.parse import urlencode
 
 import elasticapm
 from elasticapm.base import Client
@@ -44,6 +45,8 @@
 from elasticapm.utils.disttracing import TraceParent
 from elasticapm.utils.logging import get_logger
 
+SERVERLESS_HTTP_REQUEST = ("api", "elb")
+
 logger = get_logger("elasticapm.serverless")
 
 COLD_START = True
@@ -134,18 +137,25 @@ def __enter__(self):
         transaction_type = "request"
         transaction_name = os.environ.get("AWS_LAMBDA_FUNCTION_NAME", self.name)
 
-        self.httpmethod = nested_key(self.event, "requestContext", "httpMethod") or nested_key(
-            self.event, "requestContext", "http", "method"
+        self.httpmethod = (
+            nested_key(self.event, "requestContext", "httpMethod")
+            or nested_key(self.event, "requestContext", "http", "method")
+            or nested_key(self.event, "httpMethod")
         )
-        if self.httpmethod:  # API Gateway
-            self.source = "api"
-            if nested_key(self.event, "requestContext", "httpMethod"):
+
+        if self.httpmethod:  # http request
+            if nested_key(self.event, "requestContext", "elb"):
+                self.source = "elb"
+                resource = nested_key(self.event, "path")
+            elif nested_key(self.event, "requestContext", "httpMethod"):
+                self.source = "api"
                 # API v1
                 resource = "/{}{}".format(
                     nested_key(self.event, "requestContext", "stage"),
                     nested_key(self.event, "requestContext", "resourcePath"),
                 )
             else:
+                self.source = "api"
                 # API v2
                 route_key = nested_key(self.event, "requestContext", "routeKey")
                 route_key = f"/{route_key}" if route_key.startswith("$") else route_key.split(" ", 1)[-1]
@@ -170,7 +180,7 @@ def __enter__(self):
 
         self.transaction = self.client.begin_transaction(transaction_type, trace_parent=trace_parent)
         elasticapm.set_transaction_name(transaction_name, override=False)
-        if self.source == "api":
+        if self.source in SERVERLESS_HTTP_REQUEST:
             elasticapm.set_context(
                 lambda: get_data_from_request(
                     self.event,
@@ -199,7 +209,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     elasticapm.set_transaction_result("HTTP 5xx", override=False)
         if exc_val:
             self.client.capture_exception(exc_info=(exc_type, exc_val, exc_tb), handled=False)
-            if self.source == "api":
+            if self.source in SERVERLESS_HTTP_REQUEST:
                 elasticapm.set_transaction_result("HTTP 5xx", override=False)
                 elasticapm.set_transaction_outcome(http_status_code=500, override=False)
                 elasticapm.set_context({"status_code": 500}, "response")
@@ -246,6 +256,17 @@ def set_metadata_and_context(self, coldstart: bool) -> None:
             cloud_context["origin"]["service"] = {"name": "api gateway"}
             cloud_context["origin"]["account"] = {"id": self.event["requestContext"]["accountId"]}
             cloud_context["origin"]["provider"] = "aws"
+        elif self.source == "elb":
+            elb_target_group_arn = self.event["requestContext"]["elb"]["targetGroupArn"]
+            faas["trigger"]["type"] = "http"
+            faas["trigger"]["request_id"] = self.event["headers"]["x-amzn-trace-id"]
+            service_context["origin"] = {"name": elb_target_group_arn.split(":")[5]}
+            service_context["origin"]["id"] = elb_target_group_arn
+            cloud_context["origin"] = {}
+            cloud_context["origin"]["service"] = {"name": "elb"}
+            cloud_context["origin"]["account"] = {"id": elb_target_group_arn.split(":")[4]}
+            cloud_context["origin"]["region"] = elb_target_group_arn.split(":")[3]
+            cloud_context["origin"]["provider"] = "aws"
         elif self.source == "sqs":
             record = self.event["Records"][0]
             faas["trigger"]["type"] = "pubsub"
@@ -281,7 +302,7 @@ def set_metadata_and_context(self, coldstart: bool) -> None:
             service_context["origin"]["service"] = {"name": "sns"}
             cloud_context["origin"] = {}
             cloud_context["origin"]["region"] = record["Sns"]["TopicArn"].split(":")[3]
-            cloud_context["origin"]["account_id"] = record["Sns"]["TopicArn"].split(":")[4]
+            cloud_context["origin"]["account"] = {"id": record["Sns"]["TopicArn"].split(":")[4]}
             cloud_context["origin"]["provider"] = "aws"
             message_context["queue"] = {"name": service_context["origin"]["name"]}
             if "Timestamp" in record["Sns"]:
@@ -354,7 +375,13 @@ def get_data_from_request(event: dict, capture_body: bool = False, capture_heade
     result = {}
     if capture_headers and "headers" in event:
         result["headers"] = event["headers"]
-    method = nested_key(event, "requestContext", "httpMethod") or nested_key(event, "requestContext", "http", "method")
+
+    method = (
+        nested_key(event, "requestContext", "httpMethod")
+        or nested_key(event, "requestContext", "http", "method")
+        or nested_key(event, "httpMethod")
+    )
+
     if not method:
         # Not API Gateway
         return result
@@ -405,16 +432,23 @@ def get_url_dict(event: dict) -> dict:
     headers = event.get("headers", {})
     protocol = headers.get("X-Forwarded-Proto", headers.get("x-forwarded-proto", "https"))
     host = headers.get("Host", headers.get("host", ""))
-    stage = "/" + (nested_key(event, "requestContext", "stage") or "")
-    path = event.get("path", event.get("rawPath", "").split(stage)[-1])
+    stage = nested_key(event, "requestContext", "stage") or ""
+    raw_path = event.get("rawPath", "")
+    if stage:
+        stage = "/" + stage
+        raw_path = raw_path.split(stage)[-1]
+
+    path = event.get("path", raw_path)
     port = headers.get("X-Forwarded-Port", headers.get("x-forwarded-port"))
     query = ""
     if "rawQueryString" in event:
         query = event["rawQueryString"]
     elif event.get("queryStringParameters"):
-        query = "?"
-        for k, v in event["queryStringParameters"].items():
-            query += "{}={}".format(k, v)
+        if stage:  # api requires parameters encoding to build correct url
+            query = "?" + urlencode(event["queryStringParameters"])
+        else:  # for elb we do not have the stage
+            query = "?" + "&".join(["{}={}".format(k, v) for k, v in event["queryStringParameters"].items()])
+
     url = protocol + "://" + host + stage + path + query
 
     url_dict = {

tests/contrib/serverless/aws_api_test_data.json

@@ -78,7 +78,9 @@
             "https"
         ]
     },
-    "queryStringParameters": null,
+    "queryStringParameters": {
+        "test@key": "test@value"
+    },
     "multiValueQueryStringParameters": null,
     "pathParameters": null,
     "stageVariables": null,
@@ -114,4 +116,4 @@
     },
     "body": null,
     "isBase64Encoded": false
-}
+}

tests/contrib/serverless/aws_elb_test_data.json

@@ -0,0 +1,25 @@
+{
+    "requestContext": {
+        "elb": {
+            "targetGroupArn": "arn:aws:elasticloadbalancing:us-east-2:123456789012:targetgroup/lambda-279XGJDqGZ5rsrHC2Fjr/49e9d65c45c6791a"
+        }
+    },
+    "httpMethod": "POST",
+    "path": "/toolz/api/v2.0/downloadPDF/PDF_2020-09-11_11-06-01.pdf",
+    "queryStringParameters": {
+        "test%40key": "test%40value",
+        "language": "en-DE"
+    },
+    "headers": {
+        "accept-encoding": "gzip,deflate",
+        "connection": "Keep-Alive",
+        "host": "blabla.com",
+        "user-agent": "Apache-HttpClient/4.5.13 (Java/11.0.15)",
+        "x-amzn-trace-id": "Root=1-xxxxxxxxxxxxxx",
+        "x-forwarded-for": "199.99.99.999",
+        "x-forwarded-port": "443",
+        "x-forwarded-proto": "https"
+    },
+    "body": "blablablabody",
+    "isBase64Encoded": false
+}

tests/contrib/serverless/aws_tests.py

@@ -53,6 +53,13 @@ def event_api2():
         return json.load(f)
 
 
+@pytest.fixture
+def event_elb():
+    aws_data_file = os.path.join(os.path.dirname(__file__), "aws_elb_test_data.json")
+    with open(aws_data_file) as f:
+        return json.load(f)
+
+
 @pytest.fixture
 def event_s3():
     aws_data_file = os.path.join(os.path.dirname(__file__), "aws_s3_test_data.json")
@@ -100,7 +107,7 @@ def test_request_data(event_api, event_api2):
     data = get_data_from_request(event_api, capture_body=True, capture_headers=True)
 
     assert data["method"] == "GET"
-    assert data["url"]["full"] == "https://02plqthge2.execute-api.us-east-1.amazonaws.com/dev/fetch_all"
+    assert data["url"]["full"] == "https://02plqthge2.execute-api.us-east-1.amazonaws.com/dev/fetch_all?test%40key=test%40value"
     assert data["headers"]["Host"] == "02plqthge2.execute-api.us-east-1.amazonaws.com"
 
     data = get_data_from_request(event_api2, capture_body=True, capture_headers=True)
@@ -112,10 +119,28 @@ def test_request_data(event_api, event_api2):
     data = get_data_from_request(event_api, capture_body=False, capture_headers=False)
 
     assert data["method"] == "GET"
-    assert data["url"]["full"] == "https://02plqthge2.execute-api.us-east-1.amazonaws.com/dev/fetch_all"
+    assert data["url"]["full"] == "https://02plqthge2.execute-api.us-east-1.amazonaws.com/dev/fetch_all?test%40key=test%40value"
     assert "headers" not in data
 
 
+def test_elb_request_data(event_elb):
+    data = get_data_from_request(event_elb, capture_body=True, capture_headers=True)
+
+    assert data["method"] == "POST"
+    assert data["url"][
+               "full"] == "https://blabla.com/toolz/api/v2.0/downloadPDF/PDF_2020-09-11_11-06-01.pdf?test%40key=test%40value&language=en-DE"
+    assert data["headers"]["host"] == "blabla.com"
+    assert data["body"] == "blablablabody"
+
+    data = get_data_from_request(event_elb, capture_body=False, capture_headers=False)
+
+    assert data["method"] == "POST"
+    assert data["url"][
+               "full"] == "https://blabla.com/toolz/api/v2.0/downloadPDF/PDF_2020-09-11_11-06-01.pdf?test%40key=test%40value&language=en-DE"
+    assert "headers" not in data
+    assert data["body"] == "[REDACTED]"
+
+
 def test_response_data():
     response = {"statusCode": "200", "headers": {"foo": "bar"}}
     data = get_data_from_response(response, capture_headers=True)
@@ -136,7 +161,6 @@ def test_response_data():
 
 
 def test_capture_serverless_api_gateway(event_api, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -159,7 +183,6 @@ def test_func(event, context):
 
 
 def test_capture_serverless_api_gateway_v2(event_api2, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -181,8 +204,29 @@ def test_func(event, context):
     assert transaction["context"]["response"]["status_code"] == 200
 
 
-def test_capture_serverless_s3(event_s3, context, elasticapm_client):
+def test_capture_serverless_elb(event_elb, context, elasticapm_client):
+    os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
+
+    @capture_serverless(elasticapm_client=elasticapm_client)
+    def test_func(event, context):
+        with capture_span("test_span"):
+            time.sleep(0.01)
+        return {"statusCode": 200, "headers": {"foo": "bar"}}
+
+    test_func(event_elb, context)
 
+    assert len(elasticapm_client.events[constants.TRANSACTION]) == 1
+    transaction = elasticapm_client.events[constants.TRANSACTION][0]
+
+    assert transaction["name"] == "POST /toolz/api/v2.0/downloadPDF/PDF_2020-09-11_11-06-01.pdf"
+    assert transaction["result"] == "HTTP 2xx"
+    assert transaction["span_count"]["started"] == 1
+    assert transaction["context"]["request"]["method"] == "POST"
+    assert transaction["context"]["request"]["headers"]
+    assert transaction["context"]["response"]["status_code"] == 200
+
+
+def test_capture_serverless_s3(event_s3, context, elasticapm_client):
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -201,7 +245,6 @@ def test_func(event, context):
 
 
 def test_capture_serverless_sns(event_sns, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -222,7 +265,6 @@ def test_func(event, context):
 
 
 def test_capture_serverless_sqs(event_sqs, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -243,7 +285,6 @@ def test_func(event, context):
 
 
 def test_capture_serverless_s3_batch(event_s3_batch, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)
@@ -263,7 +304,6 @@ def test_func(event, context):
 
 @pytest.mark.parametrize("elasticapm_client", [{"service_name": "override"}], indirect=True)
 def test_service_name_override(event_api, context, elasticapm_client):
-
     os.environ["AWS_LAMBDA_FUNCTION_NAME"] = "test_func"
 
     @capture_serverless(elasticapm_client=elasticapm_client)