Repository: elastic/apm-agent-python

Issue:
AWS lambda: capture_serverless does not provide information when lambda is connected to ELB (not to API GW)
Event that lambda handler receiving in case of attachment through the ELB you can find in attachment: [elb-lambda-event.txt](https://github.com/elastic/apm-agent-python/files/9103558/elb-lambda-event.txt)

Please note that queryStringParameters are quoted (in the event which is coming from API gateway they are not).
Note that in this case requestContext contains only elb element and nothing else.


**Is your feature request related to a problem? Please describe.**
Recently AWS has introduced that lambda could be joined directly to ELB (without API gateway) and when transaction is getting recorded by capture_serverless then it has much less information comparing to the API gateway connected one.

This function should be rewritten to react accordingly. For the lambdas which are ELB api invocations - the " elasticapm.set_context(lambda: get_data_from_request( is not getting invoked.

    def __enter__(self):

..... 

 self.httpmethod = nested_key(self.event, "requestContext", "httpMethod") or nested_key(
            self.event, "requestContext", "http", "method"
        )
....

if self.httpmethod:  # API Gateway
    self.source = "api"

.....

.....
       if self.source == "api":
            elasticapm.set_context(
                lambda: get_data_from_request(


Discussion:
@basepi Could you take a look please :)
Thanks for providing a sample event! We can definitely get it updated to deal with these differences from API Gateway.

Resolving patch:
diff --git a/elasticapm/contrib/serverless/aws.py b/elasticapm/contrib/serverless/aws.py
--- a/elasticapm/contrib/serverless/aws.py
+++ b/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 = {


Test patch:
diff --git a/tests/contrib/serverless/aws_api_test_data.json b/tests/contrib/serverless/aws_api_test_data.json
--- a/tests/contrib/serverless/aws_api_test_data.json
+++ b/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
-}
\ No newline at end of file
+}
diff --git a/tests/contrib/serverless/aws_elb_test_data.json b/tests/contrib/serverless/aws_elb_test_data.json
new file mode 100644
--- /dev/null
+++ b/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
+}
diff --git a/tests/contrib/serverless/aws_tests.py b/tests/contrib/serverless/aws_tests.py
--- a/tests/contrib/serverless/aws_tests.py
+++ b/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)