Add minimal OpenTelemetry instrumentation (#150)

Author: pquentin

elastic_transport/_node/_http_urllib3.py

@@ -24,7 +24,7 @@
 try:
     from importlib import metadata
 except ImportError:
-    import importlib_metadata as metadata  # type: ignore[import-not-found,no-redef]
+    import importlib_metadata as metadata  # type: ignore[no-redef]
 
 import urllib3
 from urllib3.exceptions import ConnectTimeoutError, NewConnectionError, ReadTimeoutError

elastic_transport/_otel.py

@@ -0,0 +1,52 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+from __future__ import annotations
+
+import contextlib
+import os
+import typing
+
+try:
+    from opentelemetry import trace
+
+    _tracer: trace.Tracer | None = trace.get_tracer("elastic-transport")
+except ModuleNotFoundError:
+    _tracer = None
+
+
+ENABLED_ENV_VAR = "OTEL_PYTHON_INSTRUMENTATION_ELASTICSEARCH_ENABLED"
+
+
+class OpenTelemetry:
+    def __init__(self, enabled: bool | None = None, tracer: trace.Tracer | None = None):
+        if enabled is None:
+            enabled = os.environ.get(ENABLED_ENV_VAR, "false") != "false"
+        self.tracer = tracer or _tracer
+        self.enabled = enabled and self.tracer is not None
+
+    @contextlib.contextmanager
+    def span(self, method: str) -> typing.Generator[None, None, None]:
+        if not self.enabled or self.tracer is None:
+            yield
+            return
+
+        span_name = method
+        with self.tracer.start_as_current_span(span_name) as span:
+            span.set_attribute("http.request.method", method)
+            span.set_attribute("db.system", "elasticsearch")
+            yield

elastic_transport/_transport.py

@@ -60,6 +60,7 @@
     Urllib3HttpNode,
 )
 from ._node_pool import NodePool, NodeSelector
+from ._otel import OpenTelemetry
 from ._serializer import DEFAULT_SERIALIZERS, Serializer, SerializerCollection
 from ._version import __version__
 from .client_utils import client_meta_version, resolve_default
@@ -225,6 +226,9 @@ def __init__(
         self.retry_on_status = retry_on_status
         self.retry_on_timeout = retry_on_timeout
 
+        # Instrumentation
+        self.otel = OpenTelemetry()
+
         # Build the NodePool from all the options
         node_pool_kwargs: Dict[str, Any] = {}
         if node_selector_class is not None:
@@ -252,7 +256,7 @@ def __init__(
         if sniff_on_start:
             self.sniff(True)
 
-    def perform_request(  # type: ignore[return]
+    def perform_request(
         self,
         method: str,
         target: str,
@@ -289,6 +293,32 @@ def perform_request(  # type: ignore[return]
         :arg client_meta: Extra client metadata key-value pairs to send in the client meta header.
         :returns: Tuple of the :class:`elastic_transport.ApiResponseMeta` with the deserialized response.
         """
+        with self.otel.span(method):
+            return self._perform_request(
+                method,
+                target,
+                body=body,
+                headers=headers,
+                max_retries=max_retries,
+                retry_on_status=retry_on_status,
+                retry_on_timeout=retry_on_timeout,
+                request_timeout=request_timeout,
+                client_meta=client_meta,
+            )
+
+    def _perform_request(  # type: ignore[return]
+        self,
+        method: str,
+        target: str,
+        *,
+        body: Optional[Any] = None,
+        headers: Union[Mapping[str, Any], DefaultType] = DEFAULT,
+        max_retries: Union[int, DefaultType] = DEFAULT,
+        retry_on_status: Union[Collection[int], DefaultType] = DEFAULT,
+        retry_on_timeout: Union[bool, DefaultType] = DEFAULT,
+        request_timeout: Union[Optional[float], DefaultType] = DEFAULT,
+        client_meta: Union[Tuple[Tuple[str, str], ...], DefaultType] = DEFAULT,
+    ) -> TransportApiResponse:
         if headers is DEFAULT:
             request_headers = HttpHeaders()
         else:

setup.py

@@ -52,7 +52,6 @@
     install_requires=[
         "urllib3>=1.26.2, <3",
         "certifi",
-        "dataclasses; python_version<'3.7'",
         "importlib-metadata; python_version<'3.8'",
     ],
     python_requires=">=3.7",
@@ -69,6 +68,8 @@
             "aiohttp",
             "httpx",
             "respx",
+            "opentelemetry-api",
+            "opentelemetry-sdk",
             # Override Read the Docs default (sphinx<2)
             "sphinx>2",
             "furo",

tests/test_otel.py

@@ -0,0 +1,41 @@
+#  Licensed to Elasticsearch B.V. under one or more contributor
+#  license agreements. See the NOTICE file distributed with
+#  this work for additional information regarding copyright
+#  ownership. Elasticsearch B.V. licenses this file to you under
+#  the Apache License, Version 2.0 (the "License"); you may
+#  not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing,
+#  software distributed under the License is distributed on an
+#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+#  KIND, either express or implied.  See the License for the
+#  specific language governing permissions and limitations
+#  under the License.
+
+
+from opentelemetry.sdk.trace import TracerProvider, export
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from elastic_transport._otel import OpenTelemetry
+
+
+def test_span():
+    tracer_provider = TracerProvider()
+    memory_exporter = InMemorySpanExporter()
+    span_processor = export.SimpleSpanProcessor(memory_exporter)
+    tracer_provider.add_span_processor(span_processor)
+    tracer = tracer_provider.get_tracer(__name__)
+
+    otel = OpenTelemetry(enabled=True, tracer=tracer)
+    with otel.span("GET"):
+        pass
+
+    spans = memory_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].attributes == {
+        "http.request.method": "GET",
+        "db.system": "elasticsearch",
+    }