Merge branch 'main' into move-to-mocket

Author: xrmx

.changelog/4216.fixed

@@ -0,0 +1 @@
+`opentelemetry-instrumentation-aws-lambda`: fix improper handling of header casing

.changelog/4505.fixed

@@ -0,0 +1 @@
+`opentelemetry-instrumentation-celery`: add null guards and type-safe helper handling around Celery context propagation internals

docs/nitpick-exceptions.ini

@@ -32,6 +32,7 @@ py-class=
     httpx.URL
     httpx.Headers
     aiohttp.web_request.Request
+    celery.worker.request.Request
     yarl.URL
     cimpl.Producer
     cimpl.Consumer

instrumentation/opentelemetry-instrumentation-aws-lambda/src/opentelemetry/instrumentation/aws_lambda/__init__.py

@@ -73,6 +73,7 @@ def custom_event_context_extractor(lambda_event):
 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.cidict import CIDict
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.metrics import MeterProvider, get_meter_provider
@@ -165,7 +166,9 @@ def _default_event_context_extractor(lambda_event: Any) -> Context:
         )
     if not isinstance(headers, dict):
         headers = {}
-    return get_global_textmap().extract(headers)
+    return get_global_textmap().extract(
+        CIDict(headers),
+    )
 
 
 def _determine_parent_context(
@@ -205,20 +208,21 @@ def _set_api_gateway_v1_proxy_attributes(
     span.set_attribute(HTTP_METHOD, lambda_event.get("httpMethod"))
 
     if lambda_event.get("headers"):
-        if "User-Agent" in lambda_event["headers"]:
+        headers = CIDict(lambda_event["headers"])
+        if "User-Agent" in headers:
             span.set_attribute(
                 HTTP_USER_AGENT,
-                lambda_event["headers"]["User-Agent"],
+                headers["User-Agent"],
             )
-        if "X-Forwarded-Proto" in lambda_event["headers"]:
+        if "X-Forwarded-Proto" in headers:
             span.set_attribute(
                 HTTP_SCHEME,
-                lambda_event["headers"]["X-Forwarded-Proto"],
+                headers["X-Forwarded-Proto"],
             )
-        if "Host" in lambda_event["headers"]:
+        if "Host" in headers:
             span.set_attribute(
                 NET_HOST_NAME,
-                lambda_event["headers"]["Host"],
+                headers["Host"],
             )
     if "resource" in lambda_event:
         span.set_attribute(HTTP_ROUTE, lambda_event["resource"])

instrumentation/opentelemetry-instrumentation-aws-lambda/tests/test_aws_lambda_instrumentation_manual.py

@@ -314,6 +314,36 @@ def custom_event_context_extractor(lambda_event):
                 expected_baggage=MOCK_W3C_BAGGAGE_VALUE,
                 propagators="tracecontext,baggage",
             ),
+            TestCase(
+                name="case_insensitive_headers_uppercase",
+                custom_extractor=None,
+                context={
+                    "headers": {
+                        TraceContextTextMapPropagator._TRACEPARENT_HEADER_NAME.upper(): MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        TraceContextTextMapPropagator._TRACESTATE_HEADER_NAME.upper(): f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                },
+                expected_traceid=MOCK_W3C_TRACE_ID,
+                expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
+                expected_trace_state_len=3,
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+            ),
+            TestCase(
+                name="case_insensitive_headers_mixedcase",
+                custom_extractor=None,
+                context={
+                    "headers": {
+                        "TraceParent": MOCK_W3C_TRACE_CONTEXT_SAMPLED,
+                        "tRaCeStAtE": f"{MOCK_W3C_TRACE_STATE_KEY}={MOCK_W3C_TRACE_STATE_VALUE},foo=1,bar=2",
+                    }
+                },
+                expected_traceid=MOCK_W3C_TRACE_ID,
+                expected_parentid=MOCK_W3C_PARENT_SPAN_ID,
+                expected_trace_state_len=3,
+                expected_state_value=MOCK_W3C_TRACE_STATE_VALUE,
+                xray_traceid=MOCK_XRAY_TRACE_CONTEXT_NOT_SAMPLED,
+            ),
         ]
         for test in tests:
             with self.subTest(test_name=test.name):
@@ -389,6 +419,57 @@ def test_lambda_no_error_with_invalid_flush_timeout(self):
 
         test_env_patch.stop()
 
+    def test_api_gateway_v1_attributes_case_insensitivity(self):
+        AwsLambdaInstrumentor().instrument()
+
+        mock_execute_lambda(
+            {
+                "httpMethod": "GET",
+                "headers": {
+                    "user-agent": "lowercase-agent",
+                    "host": "lowercase-host",
+                    "x-forwarded-proto": "http",
+                },
+                "resource": "/test",
+                "requestContext": {
+                    "version": "1.0",
+                },
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(
+            span.attributes.get(HTTP_USER_AGENT), "lowercase-agent"
+        )
+        self.assertEqual(span.attributes.get(NET_HOST_NAME), "lowercase-host")
+        self.assertEqual(span.attributes.get(HTTP_SCHEME), "http")
+
+        self.memory_exporter.clear()
+
+        mock_execute_lambda(
+            {
+                "httpMethod": "GET",
+                "headers": {
+                    "uSeR-aGeNt": "mixed-agent",
+                    "hOsT": "mixed-host",
+                    "X-fOrWaRdEd-PrOtO": "https",
+                },
+                "resource": "/test",
+                "requestContext": {
+                    "version": "1.0",
+                },
+            }
+        )
+
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.attributes.get(HTTP_USER_AGENT), "mixed-agent")
+        self.assertEqual(span.attributes.get(NET_HOST_NAME), "mixed-host")
+        self.assertEqual(span.attributes.get(HTTP_SCHEME), "https")
+
     def test_lambda_handles_multiple_consumers(self):
         test_env_patch = mock.patch.dict(
             "os.environ",

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py

@@ -48,13 +48,16 @@ def add(x, y):
 ---
 """
 
+from __future__ import annotations
+
 import logging
+from collections.abc import Collection, Iterable
 from timeit import default_timer
-from typing import Collection, Iterable
 
 from billiard import VERSION
 from billiard.einfo import ExceptionInfo
 from celery import signals  # pylint: disable=no-name-in-module
+from celery.worker.request import Request  # pylint: disable=no-name-in-module
 
 from opentelemetry import context as context_api
 from opentelemetry import trace
@@ -88,8 +91,8 @@ def add(x, y):
 _TASK_NAME_KEY = "celery.task_name"
 
 
-class CeleryGetter(Getter):
-    def get(self, carrier, key):
+class CeleryGetter(Getter[Request]):
+    def get(self, carrier: Request, key: str) -> list[str] | None:
         value = getattr(carrier, key, None)
         if value is None:
             return None
@@ -98,16 +101,12 @@ def get(self, carrier, key):
         # of ints).  The TextMapPropagator contract requires string
         # values, so coerce anything that isn't already a string.
         if isinstance(value, str):
-            value = (value,)
-        elif isinstance(value, Iterable):
-            value = tuple(
-                str(v) if not isinstance(v, str) else v for v in value
-            )
-        else:
-            value = (str(value),)
-        return value
+            return [value]
+        if isinstance(value, Iterable):
+            return [str(v) if not isinstance(v, str) else v for v in value]
+        return [str(value)]
 
-    def keys(self, carrier):
+    def keys(self, carrier: Request) -> list[str]:
         return []
 
 

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py

@@ -4,7 +4,8 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Optional, Tuple
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Optional, Protocol, cast
 
 from celery import registry  # pylint: disable=no-name-in-module
 from celery.app.task import Task
@@ -18,6 +19,15 @@
 if TYPE_CHECKING:
     from contextlib import AbstractContextManager
 
+ContextKey = tuple[str, bool]
+ContextTuple = tuple[Span, "AbstractContextManager[Span]", object | None]
+ContextDict = dict[ContextKey, ContextTuple]
+
+
+class ContextCarrier(Protocol):
+    def get(self, key: str, default: Any = None) -> Any: ...
+
+
 logger = logging.getLogger(__name__)
 
 # Celery Context key
@@ -48,7 +58,10 @@
 
 
 # pylint:disable=too-many-branches
-def set_attributes_from_context(span, context):
+def set_attributes_from_context(
+    span: Span,
+    context: ContextCarrier,
+) -> None:
     """Helper to extract meta values from a Celery Context"""
     if not span.is_recording():
         return
@@ -144,7 +157,7 @@ def attach_context(
     if task is None:
         return
 
-    ctx_dict = getattr(task, CTX_KEY, None)
+    ctx_dict = cast(Optional[ContextDict], getattr(task, CTX_KEY, None))
 
     if ctx_dict is None:
         ctx_dict = {}
@@ -153,12 +166,17 @@ def attach_context(
     ctx_dict[(task_id, is_publish)] = (span, activation, token)
 
 
-def detach_context(task, task_id, is_publish=False) -> None:
+def detach_context(
+    task: Optional[Task], task_id: str, is_publish: bool = False
+) -> None:
     """Helper to remove  `Span`, `ContextManager` and context token in a
     Celery task when it's propagated.
     This function handles tasks where no values are attached to the `Task`.
     """
-    span_dict = getattr(task, CTX_KEY, None)
+    if task is None:
+        return
+
+    span_dict = cast(Optional[ContextDict], getattr(task, CTX_KEY, None))
     if span_dict is None:
         return
 
@@ -167,27 +185,30 @@ def detach_context(task, task_id, is_publish=False) -> None:
 
 
 def retrieve_context(
-    task, task_id, is_publish=False
-) -> Optional[Tuple[Span, AbstractContextManager[Span], Optional[object]]]:
+    task: Optional[Task], task_id: str, is_publish: bool = False
+) -> Optional[ContextTuple]:
     """Helper to retrieve an active `Span`, `ContextManager` and context token
     stored in a `Task` instance
     """
-    span_dict = getattr(task, CTX_KEY, None)
+    if task is None:
+        return None
+
+    span_dict = cast(Optional[ContextDict], getattr(task, CTX_KEY, None))
     if span_dict is None:
         return None
 
     # See note in `attach_context` for key info
     return span_dict.get((task_id, is_publish), None)
 
 
-def retrieve_task(kwargs):
+def retrieve_task(kwargs: Mapping[str, Any]) -> Optional[Task]:
     task = kwargs.get("task")
     if task is None:
         logger.debug("Unable to retrieve task from signal arguments")
-    return task
+    return cast(Optional[Task], task)
 
 
-def retrieve_task_from_sender(kwargs):
+def retrieve_task_from_sender(kwargs: Mapping[str, Any]) -> Optional[Task]:
     sender = kwargs.get("sender")
     if sender is None:
         logger.debug("Unable to retrieve the sender from signal arguments")
@@ -199,30 +220,31 @@ def retrieve_task_from_sender(kwargs):
         if sender is None:
             logger.debug("Unable to retrieve the task from sender=%s", sender)
 
-    return sender
+    return cast(Optional[Task], sender)
 
 
-def retrieve_task_id(kwargs):
+def retrieve_task_id(kwargs: Mapping[str, Any]) -> Optional[str]:
     task_id = kwargs.get("task_id")
     if task_id is None:
         logger.debug("Unable to retrieve task_id from signal arguments")
-    return task_id
+    return cast(Optional[str], task_id)
 
 
-def retrieve_task_id_from_request(kwargs):
+def retrieve_task_id_from_request(kwargs: Mapping[str, Any]) -> Optional[str]:
     # retry signal does not include task_id as argument so use request argument
     request = kwargs.get("request")
     if request is None:
         logger.debug("Unable to retrieve the request from signal arguments")
+        return None
 
-    task_id = getattr(request, "id")
+    task_id = cast(Optional[str], getattr(request, "id", None))
     if task_id is None:
         logger.debug("Unable to retrieve the task_id from the request")
 
     return task_id
 
 
-def retrieve_task_id_from_message(kwargs):
+def retrieve_task_id_from_message(kwargs: Mapping[str, Any]) -> Optional[str]:
     """Helper to retrieve the `Task` identifier from the message `body`.
     This helper supports Protocol Version 1 and 2. The Protocol is well
     detailed in the official documentation:
@@ -232,12 +254,14 @@ def retrieve_task_id_from_message(kwargs):
     body = kwargs.get("body")
     if headers is not None and len(headers) > 0:
         # Protocol Version 2 (default from Celery 4.0)
-        return headers.get("id")
+        return cast(Optional[str], headers.get("id"))
     # Protocol Version 1
-    return body.get("id")
+    if body is None:
+        return None
+    return cast(Optional[str], body.get("id"))
 
 
-def retrieve_reason(kwargs):
+def retrieve_reason(kwargs: Mapping[str, Any]) -> Optional[object]:
     reason = kwargs.get("reason")
     if not reason:
         logger.debug("Unable to retrieve the retry reason")