Record OpenTelemetry spans

pquentin · pquentin · commit ae3606ce0f52 · 2024-06-24T17:24:57.000+04:00
diff --git a/elasticsearch_serverless/_async/client/_base.py b/elasticsearch_serverless/_async/client/_base.py
@@ -27,10 +27,12 @@
     HttpHeaders,
     ListApiResponse,
     ObjectApiResponse,
+    OpenTelemetrySpan,
     TextApiResponse,
 )
 from elastic_transport.client_utils import DEFAULT, DefaultType
 
+from ..._otel import OpenTelemetry
 from ...compat import warn_stacklevel
 from ...exceptions import (
     HTTP_EXCEPTIONS,
@@ -125,6 +127,7 @@ def __init__(self, _transport: AsyncTransport) -> None:
         self._retry_on_timeout: Union[DefaultType, bool] = DEFAULT
         self._retry_on_status: Union[DefaultType, Collection[int]] = DEFAULT
         self._verified_elasticsearch = False
+        self._otel = OpenTelemetry()
 
     @property
     def transport(self) -> AsyncTransport:
@@ -140,6 +143,32 @@ async def perform_request(
         body: Optional[Any] = None,
         endpoint_id: Optional[str] = None,
         path_parts: Optional[Mapping[str, Any]] = None,
+    ) -> ApiResponse[Any]:
+        with self._otel.span(
+            method,
+            endpoint_id=endpoint_id,
+            path_parts=path_parts or {},
+        ) as otel_span:
+            response = await self._perform_request(
+                method,
+                path,
+                params=params,
+                headers=headers,
+                body=body,
+                otel_span=otel_span,
+            )
+            otel_span.set_elastic_cloud_metadata(response.meta.headers)
+            return response
+
+    async def _perform_request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: Optional[Mapping[str, Any]] = None,
+        headers: Optional[Mapping[str, str]] = None,
+        body: Optional[Any] = None,
+        otel_span: OpenTelemetrySpan,
     ) -> ApiResponse[Any]:
         if headers:
             request_headers = self._headers.copy()
@@ -162,6 +191,7 @@ async def perform_request(
             retry_on_status=self._retry_on_status,
             retry_on_timeout=self._retry_on_timeout,
             client_meta=self._client_meta,
+            otel_span=otel_span,
         )
 
         # HEAD with a 404 is returned as a normal response
diff --git a/elasticsearch_serverless/_otel.py b/elasticsearch_serverless/_otel.py
@@ -0,0 +1,92 @@
+#  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
+from typing import TYPE_CHECKING, Generator, Mapping
+
+if TYPE_CHECKING:
+    from typing import Literal
+
+try:
+    from opentelemetry import trace
+
+    _tracer: trace.Tracer | None = trace.get_tracer("elasticsearch-api")
+except ModuleNotFoundError:
+    _tracer = None
+
+from elastic_transport import OpenTelemetrySpan
+
+# Valid values for the enabled config are 'true' and 'false'. Default is 'true'.
+ENABLED_ENV_VAR = "OTEL_PYTHON_INSTRUMENTATION_ELASTICSEARCH_ENABLED"
+# Describes how to handle search queries in the request body when assigned to
+# a span attribute.
+# Valid values are 'omit' and 'raw'.
+# Default is 'omit' as 'raw' has security implications.
+BODY_STRATEGY_ENV_VAR = "OTEL_PYTHON_INSTRUMENTATION_ELASTICSEARCH_CAPTURE_SEARCH_QUERY"
+DEFAULT_BODY_STRATEGY = "omit"
+
+
+class OpenTelemetry:
+    def __init__(
+        self,
+        enabled: bool | None = None,
+        tracer: trace.Tracer | None = None,
+        # TODO import Literal at the top-level when dropping Python 3.7
+        body_strategy: 'Literal["omit", "raw"]' | 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
+
+        if body_strategy is not None:
+            self.body_strategy = body_strategy
+        else:
+            self.body_strategy = os.environ.get(
+                BODY_STRATEGY_ENV_VAR, DEFAULT_BODY_STRATEGY
+            )  # type: ignore[assignment]
+            assert self.body_strategy in ("omit", "raw")
+
+    @contextlib.contextmanager
+    def span(
+        self,
+        method: str,
+        *,
+        endpoint_id: str | None,
+        path_parts: Mapping[str, str],
+    ) -> Generator[OpenTelemetrySpan, None, None]:
+        if not self.enabled or self.tracer is None:
+            yield OpenTelemetrySpan(None)
+            return
+
+        span_name = endpoint_id or method
+        with self.tracer.start_as_current_span(span_name) as otel_span:
+            otel_span.set_attribute("http.request.method", method)
+            otel_span.set_attribute("db.system", "elasticsearch")
+            if endpoint_id is not None:
+                otel_span.set_attribute("db.operation", endpoint_id)
+            for key, value in path_parts.items():
+                otel_span.set_attribute(f"db.elasticsearch.path_parts.{key}", value)
+
+            yield OpenTelemetrySpan(
+                otel_span,
+                endpoint_id=endpoint_id,
+                body_strategy=self.body_strategy,
+            )
diff --git a/elasticsearch_serverless/_sync/client/_base.py b/elasticsearch_serverless/_sync/client/_base.py
@@ -26,11 +26,13 @@
     HttpHeaders,
     ListApiResponse,
     ObjectApiResponse,
+    OpenTelemetrySpan,
     TextApiResponse,
     Transport,
 )
 from elastic_transport.client_utils import DEFAULT, DefaultType
 
+from ..._otel import OpenTelemetry
 from ...compat import warn_stacklevel
 from ...exceptions import (
     HTTP_EXCEPTIONS,
@@ -125,6 +127,7 @@ def __init__(self, _transport: Transport) -> None:
         self._retry_on_timeout: Union[DefaultType, bool] = DEFAULT
         self._retry_on_status: Union[DefaultType, Collection[int]] = DEFAULT
         self._verified_elasticsearch = False
+        self._otel = OpenTelemetry()
 
     @property
     def transport(self) -> Transport:
@@ -140,6 +143,32 @@ def perform_request(
         body: Optional[Any] = None,
         endpoint_id: Optional[str] = None,
         path_parts: Optional[Mapping[str, Any]] = None,
+    ) -> ApiResponse[Any]:
+        with self._otel.span(
+            method,
+            endpoint_id=endpoint_id,
+            path_parts=path_parts or {},
+        ) as otel_span:
+            response = self._perform_request(
+                method,
+                path,
+                params=params,
+                headers=headers,
+                body=body,
+                otel_span=otel_span,
+            )
+            otel_span.set_elastic_cloud_metadata(response.meta.headers)
+            return response
+
+    def _perform_request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: Optional[Mapping[str, Any]] = None,
+        headers: Optional[Mapping[str, str]] = None,
+        body: Optional[Any] = None,
+        otel_span: OpenTelemetrySpan,
     ) -> ApiResponse[Any]:
         if headers:
             request_headers = self._headers.copy()
@@ -162,6 +191,7 @@ def perform_request(
             retry_on_status=self._retry_on_status,
             retry_on_timeout=self._retry_on_timeout,
             client_meta=self._client_meta,
+            otel_span=otel_span,
         )
 
         # HEAD with a 404 is returned as a normal response
diff --git a/test_elasticsearch_serverless/test_client/test_options.py b/test_elasticsearch_serverless/test_client/test_options.py
@@ -17,6 +17,7 @@
 #  under the License.
 
 import pytest
+from elastic_transport import OpenTelemetrySpan
 from elastic_transport.client_utils import DEFAULT
 
 from elasticsearch_serverless import AsyncElasticsearch, Elasticsearch
@@ -137,6 +138,7 @@ def test_options_passed_to_perform_request(self):
         assert call.pop("retry_on_timeout") is DEFAULT
         assert call.pop("retry_on_status") is DEFAULT
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -155,6 +157,7 @@ def test_options_passed_to_perform_request(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][1]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -180,6 +183,7 @@ def test_options_passed_to_perform_request(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][0]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -207,6 +211,7 @@ async def test_options_passed_to_async_perform_request(self):
         assert call.pop("retry_on_timeout") is DEFAULT
         assert call.pop("retry_on_status") is DEFAULT
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -225,6 +230,7 @@ async def test_options_passed_to_async_perform_request(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][1]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -250,6 +256,7 @@ async def test_options_passed_to_async_perform_request(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][0]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -389,6 +396,7 @@ def test_options_timeout_parameters(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][0]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -418,6 +426,7 @@ def test_options_timeout_parameters(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][0]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -442,6 +451,7 @@ def test_options_timeout_parameters(self):
         assert call.pop("retry_on_timeout") is DEFAULT
         assert call.pop("retry_on_status") is DEFAULT
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
@@ -463,6 +473,7 @@ def test_options_timeout_parameters(self):
         calls = client.transport.calls
         call = calls[("GET", "/test")][0]
         assert call.pop("client_meta") is DEFAULT
+        assert isinstance(call.pop("otel_span"), OpenTelemetrySpan)
         assert call == {
             "headers": {
                 "accept": "application/json",
diff --git a/test_elasticsearch_serverless/test_otel.py b/test_elasticsearch_serverless/test_otel.py
