From f7ace0ce67bffa87bea71a2f0738c88bd1d073ae Mon Sep 17 00:00:00 2001 From: Michael Lorenzana Date: Wed, 31 Jul 2024 13:34:55 -0400 Subject: [PATCH 1/2] added in sqs trace parent collection --- .../instrumentation/aws_lambda/__init__.py | 46 +++++++++++++------ 1 file changed, 33 insertions(+), 13 deletions(-) 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 index 883296d85b..e124525c11 100644 --- 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 @@ -71,7 +71,7 @@ def custom_event_context_extractor(lambda_event): import os import time from importlib import import_module -from typing import Any, Callable, Collection +from typing import Any, Callable, Collection, List, Mapping, Optional from urllib.parse import urlencode from wrapt import wrap_function_wrapper @@ -83,6 +83,7 @@ def custom_event_context_extractor(lambda_event): from opentelemetry.instrumentation.utils import unwrap from opentelemetry.metrics import MeterProvider, get_meter_provider from opentelemetry.propagate import get_global_textmap +from opentelemetry.propagators.textmap import CarrierT, Getter from opentelemetry.semconv.resource import ResourceAttributes from opentelemetry.semconv.trace import SpanAttributes from opentelemetry.trace import ( @@ -104,6 +105,20 @@ def custom_event_context_extractor(lambda_event): ) +class LambdaSqsGetter(Getter[CarrierT]): + def get(self, carrier: CarrierT, key: str) -> Optional[List[str]]: + msg_attr = carrier.get(key) + if not isinstance(msg_attr, Mapping): + return None + value = msg_attr.get("stringValue") + if value is None: + return None + return [value] + + def keys(self, carrier: CarrierT) -> List[str]: + return list(carrier.keys()) + + def _default_event_context_extractor(lambda_event: Any) -> Context: """Default way of extracting the context from the Lambda Event. @@ -117,21 +132,26 @@ def _default_event_context_extractor(lambda_event: Any) -> Context: https://docs.aws.amazon.com/apigateway/latest/developerguide/set-up-lambda-proxy-integrations.html#api-gateway-simple-proxy-for-lambda-input-format Args: - lambda_event: user-defined, so it could be anything, but this - method counts on it being a map with a 'headers' key + lambda_event: user-defined, so it could be anything, this method + will extract the context from the 'headers' key or the + 'messageAttributes' key if the record count is 1 Returns: A Context with configuration found in the event. """ - headers = None - try: - headers = lambda_event["headers"] - except (TypeError, KeyError): - logger.debug( - "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context." - ) - if not isinstance(headers, dict): - headers = {} - return get_global_textmap().extract(headers) + headers = lambda_event.get("headers") + if headers and isinstance(headers, dict): + return get_global_textmap().extract(headers) + + records = lambda_event.get("Records") + if records and isinstance(records, list) and len(records) == 1: + message_attributes = records[0]['messageAttributes'] + if message_attributes and isinstance(message_attributes, dict): + return get_global_textmap().extract(message_attributes, getter=LambdaSqsGetter()) + + logger.debug( + "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context." + ) + return get_global_textmap().extract({}) def _determine_parent_context( From 9539e7c151f5c8a111b1b1aff10e68b312fb14f7 Mon Sep 17 00:00:00 2001 From: Michael Lorenzana Date: Wed, 31 Jul 2024 14:54:47 -0400 Subject: [PATCH 2/2] added test, fixed extractor for existing tests --- .../instrumentation/aws_lambda/__init__.py | 28 ++++++++++--------- .../test_aws_lambda_instrumentation_manual.py | 14 ++++++++++ 2 files changed, 29 insertions(+), 13 deletions(-) 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 index e124525c11..fd61b97200 100644 --- 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 @@ -138,19 +138,21 @@ def _default_event_context_extractor(lambda_event: Any) -> Context: Returns: A Context with configuration found in the event. """ - headers = lambda_event.get("headers") - if headers and isinstance(headers, dict): - return get_global_textmap().extract(headers) - - records = lambda_event.get("Records") - if records and isinstance(records, list) and len(records) == 1: - message_attributes = records[0]['messageAttributes'] - if message_attributes and isinstance(message_attributes, dict): - return get_global_textmap().extract(message_attributes, getter=LambdaSqsGetter()) - - logger.debug( - "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context." - ) + try: + headers = lambda_event.get("headers") + if headers and isinstance(headers, dict): + return get_global_textmap().extract(headers) + + records = lambda_event.get("Records") + if records and isinstance(records, list) and len(records) == 1: + message_attributes = records[0]['messageAttributes'] + if message_attributes and isinstance(message_attributes, dict): + return get_global_textmap().extract(message_attributes, getter=LambdaSqsGetter()) + except (TypeError, KeyError, AttributeError): + logger.debug( + "Extracting context from Lambda Event failed: either enable X-Ray active tracing or configure API Gateway to trigger this Lambda function as a pure proxy. Otherwise, generated spans will have an invalid (empty) parent context." + ) + return get_global_textmap().extract({}) 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 index 00940547ea..5d9b21daff 100644 --- 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 @@ -266,6 +266,20 @@ def custom_event_context_extractor(lambda_event): expected_state_value=MOCK_W3C_TRACE_STATE_VALUE, xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED, ), + TestCase( + name="sqs_single_record", + custom_extractor=None, + context={ + "Records": [ + {'messageAttributes': { + TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME: {'stringValue': MOCK_W3C_TRACE_CONTEXT_SAMPLED, 'stringListValues': [], 'binaryListValues': [], 'dataType': 'String'} + }} + ] + }, + expected_traceid=MOCK_W3C_TRACE_ID, + expected_parentid=MOCK_W3C_PARENT_SPAN_ID, + xray_traceid=MOCK_XRAY_TRACE_CONTEXT_SAMPLED, + ) ] for test in tests: