fix: forward OTEL context to subthreads in parallel bulk calls

Author: Alisson Claudino

elasticsearch/_otel.py

@@ -25,7 +25,11 @@
     from typing import Literal
 
 try:
+    from opentelemetry import context as otel_context
     from opentelemetry import trace
+    from opentelemetry.trace.propagation.tracecontext import (
+        TraceContextTextMapPropagator,
+    )
 
     _tracer: trace.Tracer | None = trace.get_tracer("elasticsearch-api")
 except ModuleNotFoundError:
@@ -44,6 +48,8 @@
 
 
 class OpenTelemetry:
+    current_context = None
+
     def __init__(
         self,
         enabled: bool | None = None,
@@ -78,6 +84,8 @@ def span(
 
         span_name = endpoint_id or method
         with self.tracer.start_as_current_span(span_name) as otel_span:
+            self.current_context = {}
+            TraceContextTextMapPropagator().inject(self.current_context)
             otel_span.set_attribute("http.request.method", method)
             otel_span.set_attribute("db.system", "elasticsearch")
             if endpoint_id is not None:
@@ -90,3 +98,16 @@ def span(
                 endpoint_id=endpoint_id,
                 body_strategy=self.body_strategy,
             )
+
+    @contextlib.contextmanager
+    def recover_parent_context(self):
+        token = None
+        if self.current_context:
+            otel_parent_ctx = TraceContextTextMapPropagator().extract(
+                carrier=self.current_context
+            )
+            token = otel_context.attach(otel_parent_ctx)
+        yield
+
+        if token:
+            otel_context.detach(token)

elasticsearch/_version.py

@@ -15,4 +15,4 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-__versionstr__ = "8.14.0"
+__versionstr__ = "8.14.1-alpha"

elasticsearch/helpers/actions.py

@@ -331,28 +331,29 @@ def _process_bulk_chunk(
     """
     Send a bulk request to elasticsearch and process the output.
     """
-    if isinstance(ignore_status, int):
-        ignore_status = (ignore_status,)
-
-    try:
-        # send the actual request
-        resp = client.bulk(*args, operations=bulk_actions, **kwargs)  # type: ignore[arg-type]
-    except ApiError as e:
-        gen = _process_bulk_chunk_error(
-            error=e,
-            bulk_data=bulk_data,
-            ignore_status=ignore_status,
-            raise_on_exception=raise_on_exception,
-            raise_on_error=raise_on_error,
-        )
-    else:
-        gen = _process_bulk_chunk_success(
-            resp=resp.body,
-            bulk_data=bulk_data,
-            ignore_status=ignore_status,
-            raise_on_error=raise_on_error,
-        )
-    yield from gen
+    with client._otel.recover_parent_context():
+        if isinstance(ignore_status, int):
+            ignore_status = (ignore_status,)
+
+        try:
+            # send the actual request
+            resp = client.bulk(*args, operations=bulk_actions, **kwargs)  # type: ignore[arg-type]
+        except ApiError as e:
+            gen = _process_bulk_chunk_error(
+                error=e,
+                bulk_data=bulk_data,
+                ignore_status=ignore_status,
+                raise_on_exception=raise_on_exception,
+                raise_on_error=raise_on_error,
+            )
+        else:
+            gen = _process_bulk_chunk_success(
+                resp=resp.body,
+                bulk_data=bulk_data,
+                ignore_status=ignore_status,
+                raise_on_error=raise_on_error,
+            )
+        yield from gen
 
 
 def streaming_bulk(