Support other propagators in Python Layer

python/src/otel/otel_sdk/opentelemetry/instrumentation/aws_lambda/__init__.py

@@ -47,19 +47,27 @@ def lambda_handler(event, context):
 import logging
 import os
 from importlib import import_module
-from typing import Collection
-from wrapt import wrap_function_wrapper
+from typing import Any, Collection
 
-# TODO: aws propagator
-from opentelemetry.sdk.extension.aws.trace.propagation.aws_xray_format import (
-    AwsXRayFormat,
-)
+from opentelemetry.context.context import Context
 from opentelemetry.instrumentation.aws_lambda.package import _instruments
 from opentelemetry.instrumentation.aws_lambda.version import __version__
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.propagate import get_global_textmap
+from opentelemetry.sdk.extension.aws.trace.propagation.aws_xray_format import (
+    TRACE_HEADER_KEY,
+    AwsXRayFormat,
+)
 from opentelemetry.semconv.trace import SpanAttributes
-from opentelemetry.trace import SpanKind, get_tracer, get_tracer_provider
+from opentelemetry.trace import (
+    SpanKind,
+    Tracer,
+    get_tracer,
+    get_tracer_provider,
+)
+from opentelemetry.trace.propagation import get_current_span
+from wrapt import wrap_function_wrapper
 
 logger = logging.getLogger(__name__)
 
@@ -69,15 +77,22 @@ def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        """Instruments Lambda Handlers on AWS Lambda
+        """Instruments Lambda Handlers on AWS Lambda.
+
+        See more:
+        https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#instrumenting-aws-lambda
 
         Args:
             **kwargs: Optional arguments
                 ``tracer_provider``: a TracerProvider, defaults to global
+                ``event_context_extractor``: a method which takes the Lambda
+                    Event as input and extracts an OTel Context from it. Usually
+                    the the context is extract from HTTP headers on the event.
         """
         tracer = get_tracer(
             __name__, __version__, kwargs.get("tracer_provider")
         )
+        event_context_extractor = kwargs.get("event_context_extractor")
 
         lambda_handler = os.environ.get(
             "ORIG_HANDLER", os.environ.get("_HANDLER")
@@ -87,7 +102,10 @@ def _instrument(self, **kwargs):
         self._wrapped_function_name = wrapped_names[1]
 
         _instrument(
-            tracer, self._wrapped_module_name, self._wrapped_function_name
+            tracer,
+            self._wrapped_module_name,
+            self._wrapped_function_name,
+            event_context_extractor,
         )
 
     def _uninstrument(self, **kwargs):
@@ -97,16 +115,80 @@ def _uninstrument(self, **kwargs):
         )
 
 
-def _instrument(tracer, wrapped_module_name, wrapped_function_name):
+def _default_event_context_extractor(lambda_event: Any) -> Context:
+    """Default way of extracting the context from the Lambda Event.
+
+    Assumes the Lambda Event is a map with the headers under the 'headers' key.
+    This is the mapping to use when the Lambda is invoked by an API Gateway
+    REST API where API Gateway is acting as a pure proxy for the request.
+
+    See more:
+    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 it being a map with a 'headers' key
+    Returns:
+        A Context with configuration found in the carrier.
+    """
+    try:
+        headers = lambda_event["headers"]
+    except (TypeError, KeyError):
+        logger.warning("Failed to extract context from Lambda Event.")
+        headers = {}
+    return get_global_textmap().extract(headers)
+
+
+def _instrument(
+    tracer: Tracer,
+    wrapped_module_name,
+    wrapped_function_name,
+    event_context_extractor=None,
+):
+    def _determine_parent_context(lambda_event: Any) -> Context:
+        """Determine the parent context for the current Lambda invocation.
+
+        See more:
+        https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/trace/semantic_conventions/instrumentation/aws-lambda.md#determining-the-parent-of-a-span
+
+        Args:
+            lambda_event: user-defined, so it could be anything, but this
+                method counts it being a map with a 'headers' key
+        Returns:
+            A Context with configuration found in the carrier.
+        """
+        parent_context = None
+
+        xray_env_var = os.environ.get("_X_AMZN_TRACE_ID")
+
+        if xray_env_var:
+            parent_context = AwsXRayFormat().extract(
+                {TRACE_HEADER_KEY: xray_env_var}
+            )
+
+        if (
+            parent_context
+            and get_current_span(parent_context)
+            .get_span_context()
+            .trace_flags.sampled
+        ):
+            return parent_context
+
+        if event_context_extractor:
+            parent_context = event_context_extractor(lambda_event)
+        else:
+            parent_context = _default_event_context_extractor(lambda_event)
+
+        return parent_context
+
     def _instrumented_lambda_handler_call(call_wrapped, instance, args, kwargs):
         orig_handler_name = ".".join(
             [wrapped_module_name, wrapped_function_name]
         )
 
-        # TODO: enable propagate from AWS by env variable
-        xray_trace_id = os.environ.get("_X_AMZN_TRACE_ID", "")
-        propagator = AwsXRayFormat()
-        parent_context = propagator.extract({"X-Amzn-Trace-Id": xray_trace_id})
+        lambda_event = args[0]
+
+        parent_context = _determine_parent_context(lambda_event)
 
         with tracer.start_as_current_span(
             name=orig_handler_name, context=parent_context, kind=SpanKind.SERVER

python/src/otel/tests/test_otel.py

@@ -17,15 +17,18 @@
 from unittest import mock
 
 from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
+from opentelemetry.propagate import get_global_textmap
 from opentelemetry.sdk.extension.aws.trace.propagation.aws_xray_format import (
     TRACE_ID_FIRST_PART_LENGTH,
+    TRACE_ID_VERSION,
 )
 from opentelemetry.semconv.resource import ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.trace import SpanKind
 
 _HANDLER = "_HANDLER"
+_X_AMZN_TRACE_ID = "_X_AMZN_TRACE_ID"
 AWS_LAMBDA_EXEC_WRAPPER = "AWS_LAMBDA_EXEC_WRAPPER"
 INSTRUMENTATION_SRC_DIR = os.path.join(
     *(os.path.dirname(__file__), "..", "otel_sdk")
@@ -42,11 +45,26 @@ def __init__(self, aws_request_id, invoked_function_arn):
     aws_request_id="mock_aws_request_id",
     invoked_function_arn="arn://mock-lambda-function-arn",
 )
-MOCK_TRACE_ID = 0x5FB7331105E8BB83207FA31D4D9CDB4C
-MOCK_TRACE_ID_HEX_STR = f"{MOCK_TRACE_ID:32x}"
-MOCK_PARENT_SPAN_ID = 0x3328B8445A6DBAD2
-MOCK_PARENT_SPAN_ID_STR = f"{MOCK_PARENT_SPAN_ID:32x}"
-MOCK_LAMBDA_TRACE_CONTEXT_SAMPLED = f"Root=1-{MOCK_TRACE_ID_HEX_STR[:TRACE_ID_FIRST_PART_LENGTH]}-{MOCK_TRACE_ID_HEX_STR[TRACE_ID_FIRST_PART_LENGTH:]};Parent={MOCK_PARENT_SPAN_ID_STR};Sampled=1"
+
+MOCK_XRAY_TRACE_ID = 0x5FB7331105E8BB83207FA31D4D9CDB4C
+MOCK_XRAY_TRACE_ID_STR = f"{MOCK_XRAY_TRACE_ID:x}"
+MOCK_XRAY_PARENT_SPAN_ID = 0x3328B8445A6DBAD2
+MOCK_XRAY_TRACE_CONTEXT_COMMON = f"Root={TRACE_ID_VERSION}-{MOCK_XRAY_TRACE_ID_STR[:TRACE_ID_FIRST_PART_LENGTH]}-{MOCK_XRAY_TRACE_ID_STR[TRACE_ID_FIRST_PART_LENGTH:]};Parent={MOCK_XRAY_PARENT_SPAN_ID:x}"
+MOCK_XRAY_TRACE_CONTEXT_SAMPLED = f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=1"
+MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED = (
+    f"{MOCK_XRAY_TRACE_CONTEXT_COMMON};Sampled=0"
+)
+
+# Read more:
+# https://www.w3.org/TR/trace-context/#examples-of-http-traceparent-headers
+MOCK_W3C_TRACE_ID = 0x5CE0E9A56015FEC5AADFA328AE398115
+MOCK_W3C_PARENT_SPAN_ID = 0xAB54A98CEB1F0AD2
+MOCK_W3C_TRACE_CONTEXT_SAMPLED = (
+    f"00-{MOCK_W3C_TRACE_ID:x}-{MOCK_W3C_PARENT_SPAN_ID:x}-01"
+)
+
+MOCK_W3C_TRACE_STATE_KEY = "vendor_specific_key"
+MOCK_W3C_TRACE_STATE_VALUE = "test_value"
 
 
 def mock_aws_lambda_exec_wrapper():
@@ -63,13 +81,20 @@ def mock_aws_lambda_exec_wrapper():
     exec(open(os.path.join(INSTRUMENTATION_SRC_DIR, "otel-instrument")).read())
 
 
-def mock_execute_lambda():
+def mock_execute_lambda(event=None):
+    """Mocks Lambda importing and then calling the method at the current
+    `_HANDLER` environment variable. Like the real Lambda, if
+    `AWS_LAMBDA_EXEC_WRAPPER` is defined, if executes that before `_HANDLER`.
+
+    See more:
+    https://aws-otel.github.io/docs/getting-started/lambda/lambda-python
+    """
     if os.environ[AWS_LAMBDA_EXEC_WRAPPER]:
         globals()[os.environ[AWS_LAMBDA_EXEC_WRAPPER]]()
 
     module_name, handler_name = os.environ[_HANDLER].split(".")
     handler_module = import_module(".".join(module_name.split("/")))
-    getattr(handler_module, handler_name)("mock_event", MOCK_LAMBDA_CONTEXT)
+    getattr(handler_module, handler_name)(event, MOCK_LAMBDA_CONTEXT)
 
 
 class TestAwsLambdaInstrumentor(TestBase):
@@ -109,7 +134,7 @@ def test_active_tracing(self):
             "os.environ",
             {
                 **os.environ,
-                "_X_AMZN_TRACE_ID": MOCK_LAMBDA_TRACE_CONTEXT_SAMPLED,
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_SAMPLED,
             },
         )
         test_env_patch.start()
@@ -123,7 +148,7 @@ def test_active_tracing(self):
         self.assertEqual(len(spans), 1)
         span = spans[0]
         self.assertEqual(span.name, os.environ["ORIG_HANDLER"])
-        self.assertEqual(span.get_span_context().trace_id, MOCK_TRACE_ID)
+        self.assertEqual(span.get_span_context().trace_id, MOCK_XRAY_TRACE_ID)
         self.assertEqual(span.kind, SpanKind.SERVER)
         self.assertSpanHasAttributes(
             span,
@@ -149,7 +174,112 @@ def test_active_tracing(self):
         self.assertEqual(
             parent_context.trace_id, span.get_span_context().trace_id
         )
-        self.assertEqual(parent_context.span_id, MOCK_PARENT_SPAN_ID)
+        self.assertEqual(parent_context.span_id, MOCK_XRAY_PARENT_SPAN_ID)
+        self.assertTrue(parent_context.is_remote)
+
+        test_env_patch.stop()
+
+    def test_parent_context_from_lambda_event(self):
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                # NOT using the X-Ray Propagator
+                "OTEL_PROPAGATORS": "tracecontext",
+            },
+        )
+        test_env_patch.start()
+
+        mock_execute_lambda(
+            {
+                "headers": {
+                    "traceparent": MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                    "tracestate": f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                }
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+        parent_context = span.parent
+        self.assertEqual(
+            parent_context.trace_id, span.get_span_context().trace_id
+        )
+        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+        self.assertEqual(len(parent_context.trace_state), 3)
+        self.assertEqual(
+            parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
+            MOCK_W3C_TRACE_STATE_VALUE,
+        )
+        self.assertTrue(parent_context.is_remote)
+
+        test_env_patch.stop()
+
+    def test_using_custom_extractor(self):
+        def custom_event_context_extractor(lambda_event):
+            return get_global_textmap().extract(lambda_event["foo"]["headers"])
+
+        test_env_patch = mock.patch.dict(
+            "os.environ",
+            {
+                **os.environ,
+                # DO NOT use `otel-instrument` script, resort to "manual"
+                # instrumentation below
+                AWS_LAMBDA_EXEC_WRAPPER: "",
+                # NOT Active Tracing
+                _X_AMZN_TRACE_ID: MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+                # NOT using the X-Ray Propagator
+                "OTEL_PROPAGATORS": "tracecontext",
+            },
+        )
+        test_env_patch.start()
+
+        # NOTE: Instead of using `AWS_LAMBDA_EXEC_WRAPPER` to point `_HANDLER`
+        # to a module which instruments and calls the user `ORIG_HANDLER`, we
+        # leave `_HANDLER` as is and replace `AWS_LAMBDA_EXEC_WRAPPER` with this
+        # line below. This is like "manual" instrumentation for Lambda.
+        AwsLambdaInstrumentor().instrument(
+            event_context_extractor=custom_event_context_extractor,
+            skip_dep_check=True,
+        )
+
+        mock_execute_lambda(
+            {
+                "foo": {
+                    "headers": {
+                        "traceparent": MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        "tracestate": f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                }
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+
+        assert spans
+
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.get_span_context().trace_id, MOCK_W3C_TRACE_ID)
+
+        parent_context = span.parent
+        self.assertEqual(
+            parent_context.trace_id, span.get_span_context().trace_id
+        )
+        self.assertEqual(parent_context.span_id, MOCK_W3C_PARENT_SPAN_ID)
+        self.assertEqual(len(parent_context.trace_state), 3)
+        self.assertEqual(
+            parent_context.trace_state.get(MOCK_W3C_TRACE_STATE_KEY),
+            MOCK_W3C_TRACE_STATE_VALUE,
+        )
         self.assertTrue(parent_context.is_remote)
 
         test_env_patch.stop()