1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
# 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.
import os
from unittest import mock
import pytest
from elasticsearch import Elasticsearch, helpers
try:
from opentelemetry.sdk.trace import TracerProvider, export
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
InMemorySpanExporter,
)
except ModuleNotFoundError:
pass
from elasticsearch._otel import ENABLED_ENV_VAR, OpenTelemetry
pytestmark = [
pytest.mark.skipif(
"TEST_WITH_OTEL" not in os.environ, reason="TEST_WITH_OTEL is not set"
),
pytest.mark.otel,
]
def setup_tracing():
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__)
return tracer, memory_exporter
def test_enabled():
otel = OpenTelemetry()
assert otel.enabled == (os.environ.get(ENABLED_ENV_VAR, "true") == "true")
def test_minimal_span():
tracer, memory_exporter = setup_tracing()
otel = OpenTelemetry(enabled=True, tracer=tracer)
with otel.span("GET", endpoint_id=None, path_parts={}):
pass
spans = memory_exporter.get_finished_spans()
assert len(spans) == 1
assert spans[0].name == "GET"
assert spans[0].attributes == {
"http.request.method": "GET",
"db.system.name": "elasticsearch",
}
def test_detailed_span():
tracer, memory_exporter = setup_tracing()
otel = OpenTelemetry(enabled=True, tracer=tracer)
with otel.span(
"GET",
endpoint_id="ml.open_job",
path_parts={"job_id": "my-job"},
) as span:
span.set_elastic_cloud_metadata(
{
"X-Found-Handling-Cluster": "e9106fc68e3044f0b1475b04bf4ffd5f",
"X-Found-Handling-Instance": "instance-0000000001",
}
)
spans = memory_exporter.get_finished_spans()
assert len(spans) == 1
assert spans[0].name == "ml.open_job"
assert spans[0].attributes == {
"http.request.method": "GET",
"db.system.name": "elasticsearch",
"db.operation.name": "ml.open_job",
"db.operation.parameter.job_id": "my-job",
"db.namespace": "e9106fc68e3044f0b1475b04bf4ffd5f",
"elasticsearch.node.name": "instance-0000000001",
}
@mock.patch("elasticsearch._otel.OpenTelemetry.use_span")
@mock.patch("elasticsearch._otel.OpenTelemetry.helpers_span")
@mock.patch("elasticsearch.helpers.actions._process_bulk_chunk_success")
@mock.patch("elasticsearch.Elasticsearch.bulk")
def test_forward_otel_context_to_subthreads(
_call_bulk_mock,
_process_bulk_success_mock,
_mock_otel_helpers_span,
_mock_otel_use_span,
):
tracer, memory_exporter = setup_tracing()
es_client = Elasticsearch("http://localhost:9200")
es_client._otel = OpenTelemetry(enabled=True, tracer=tracer)
_call_bulk_mock.return_value = mock.Mock()
actions = ({"x": i} for i in range(100))
list(helpers.parallel_bulk(es_client, actions, chunk_size=4))
# Ensures that the OTEL context has been forwarded to all chunks
assert es_client._otel.helpers_span.call_count == 1
assert es_client._otel.use_span.call_count == 25
|
