Files

644 lines
22 KiB
Python

# Copyright The OpenTelemetry Authors
#
# Licensed 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.
# pylint: disable=unexpected-keyword-arg,missing-kwoa,no-value-for-parameter
import json
import os
import threading
from ast import literal_eval
from unittest import mock
import elasticsearch
import elasticsearch.exceptions
from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from pytest import mark
import opentelemetry.instrumentation.elasticsearch
from opentelemetry import trace
from opentelemetry.instrumentation.elasticsearch import (
ElasticsearchInstrumentor,
)
from opentelemetry.instrumentation.elasticsearch.utils import sanitize_body
from opentelemetry.semconv._incubating.attributes.db_attributes import (
DB_STATEMENT,
DB_SYSTEM,
)
from opentelemetry.test.test_base import TestBase
from opentelemetry.trace import StatusCode
from . import sanitization_queries # pylint: disable=no-name-in-module
major_version, minor_version = elasticsearch.VERSION[:2]
if major_version == 8:
from . import helpers_es8 as helpers # pylint: disable=no-name-in-module
elif major_version == 7:
from . import helpers_es7 as helpers # pylint: disable=no-name-in-module
elif major_version == 6:
from . import helpers_es6 as helpers # pylint: disable=no-name-in-module
Article = helpers.Article
# pylint: disable=too-many-public-methods
def normalize_arguments(doc_type, body=None):
if major_version < 7:
return (
{"body": body, "doc_type": doc_type}
if body
else {"doc_type": doc_type}
)
return {"document": body} if body else {}
def get_elasticsearch_client(*args, **kwargs):
client = Elasticsearch(*args, **kwargs)
if major_version == 8:
client._verified_elasticsearch = True
elif major_version == 7:
client.transport._verified_elasticsearch = True
return client
@mock.patch(helpers.perform_request_mock_path)
@mock.patch.dict(
os.environ, {"OTEL_PYTHON_INSTRUMENTATION_ELASTICSEARCH_ENABLED": "false"}
)
class TestElasticsearchIntegration(TestBase):
search_attributes = {
DB_SYSTEM: "elasticsearch",
"elasticsearch.url": "/test-index/_search",
"elasticsearch.method": helpers.dsl_search_method,
"elasticsearch.target": "test-index",
DB_STATEMENT: str({"query": {"bool": {"filter": "?"}}}),
}
create_attributes = {
DB_SYSTEM: "elasticsearch",
"elasticsearch.url": "/test-index",
"elasticsearch.method": "HEAD",
}
def setUp(self):
super().setUp()
self.tracer = self.tracer_provider.get_tracer(__name__)
ElasticsearchInstrumentor().instrument()
def tearDown(self):
super().tearDown()
with self.disable_logging():
ElasticsearchInstrumentor().uninstrument()
def test_instrumentor(self, request_mock):
request_mock.return_value = helpers.mock_response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.index(
index="sw",
id=1,
**normalize_arguments(body={"name": "adam"}, doc_type="_doc"),
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
# Check version and name in span's instrumentation info
self.assertEqualSpanInstrumentationScope(
span, opentelemetry.instrumentation.elasticsearch
)
# check that no spans are generated after uninstrument
ElasticsearchInstrumentor().uninstrument()
es.index(
index="sw",
id=1,
**normalize_arguments(body={"name": "adam"}, doc_type="_doc"),
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 1)
def test_span_not_recording(self, request_mock):
request_mock.return_value = (1, {}, "{}")
mock_tracer = mock.Mock()
mock_span = mock.Mock()
mock_span.is_recording.return_value = False
mock_tracer.start_span.return_value = mock_span
with mock.patch("opentelemetry.trace.get_tracer") as tracer:
tracer.return_value = mock_tracer
get_elasticsearch_client(hosts=["http://localhost:9200"])
self.assertFalse(mock_span.is_recording())
self.assertTrue(mock_span.is_recording.called)
self.assertFalse(mock_span.set_attribute.called)
self.assertFalse(mock_span.set_status.called)
ElasticsearchInstrumentor().uninstrument()
def test_prefix_arg(self, request_mock):
prefix = "prefix-from-env"
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor(span_name_prefix=prefix).instrument()
request_mock.return_value = helpers.mock_response("{}")
self._test_prefix(prefix)
def test_prefix_env(self, request_mock):
prefix = "prefix-from-args"
env_var = "OTEL_PYTHON_ELASTICSEARCH_NAME_PREFIX"
os.environ[env_var] = prefix
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument()
request_mock.return_value = helpers.mock_response("{}")
del os.environ[env_var]
self._test_prefix(prefix)
def _test_prefix(self, prefix):
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.index(
index="sw",
id=1,
**normalize_arguments(body={"name": "adam"}, doc_type="_doc"),
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
self.assertTrue(span.name.startswith(prefix))
def test_result_values(self, request_mock):
request_mock.return_value = helpers.mock_response(
'{"found": false, "timed_out": true, "took": 7}'
)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.get(
index="test-index", **normalize_arguments(doc_type="_doc"), id=1
)
spans = self.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual(spans[0].name, "Elasticsearch/test-index/_doc/:id")
self.assertEqual("False", spans[0].attributes["elasticsearch.found"])
self.assertEqual(
"True", spans[0].attributes["elasticsearch.timed_out"]
)
self.assertEqual("7", spans[0].attributes["elasticsearch.took"])
def test_trace_error_unknown(self, request_mock):
exc = RuntimeError("custom error")
request_mock.side_effect = exc
self._test_trace_error(StatusCode.ERROR, exc)
def test_trace_error_not_found(self, request_mock):
msg = "record not found"
if major_version == 8:
error = {"error": msg}
response = helpers.mock_response(
json.dumps(error), status_code=404
)
request_mock.return_value = response
exc = elasticsearch.exceptions.NotFoundError(
msg, meta=response.meta, body=None
)
else:
exc = elasticsearch.exceptions.NotFoundError(404, msg)
request_mock.side_effect = exc
self._test_trace_error(StatusCode.ERROR, exc)
def _test_trace_error(self, code, exc):
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
try:
es.get(
index="test-index",
**normalize_arguments(doc_type="_doc"),
id=1,
)
except Exception: # pylint: disable=broad-except
pass
spans = self.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertFalse(span.status.is_ok)
self.assertEqual(span.status.status_code, code)
message = getattr(exc, "message", str(exc))
self.assertEqual(
span.status.description, f"{type(exc).__name__}: {message}"
)
def test_parent(self, request_mock):
request_mock.return_value = helpers.mock_response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
with self.tracer.start_as_current_span("parent"):
es.index(
index="sw",
**normalize_arguments(doc_type="_doc", body={"name": "adam"}),
id=1,
)
spans = self.get_finished_spans()
self.assertEqual(len(spans), 2)
parent = spans.by_name("parent")
child = spans.by_name("Elasticsearch/sw/_doc/:id")
self.assertIsNotNone(child.parent)
self.assertEqual(child.parent.span_id, parent.context.span_id)
def test_multithread(self, request_mock):
request_mock.return_value = helpers.mock_response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
ev = threading.Event()
# 1. Start tracing from thread-1; make thread-2 wait
# 2. Trace something from thread-2, make thread-1 join before finishing.
# 3. Check the spans got different parents, and are in the expected order.
def target1(parent_span):
with trace.use_span(parent_span):
es.get(
index="test-index",
**normalize_arguments(doc_type="_doc"),
id=1,
)
ev.set()
ev.wait()
def target2():
ev.wait()
es.get(
index="test-index",
**normalize_arguments(doc_type="_doc"),
id=2,
)
ev.set()
with self.tracer.start_as_current_span("parent") as span:
t1 = threading.Thread(target=target1, args=(span,))
t1.start()
t2 = threading.Thread(target=target2)
t2.start()
t1.join()
t2.join()
spans = self.get_finished_spans()
self.assertEqual(3, len(spans))
s1 = spans.by_name("parent")
s2 = spans.by_attr("elasticsearch.id", "1")
s3 = spans.by_attr("elasticsearch.id", "2")
self.assertIsNotNone(s2.parent)
self.assertEqual(s2.parent.span_id, s1.context.span_id)
self.assertIsNone(s3.parent)
def test_dsl_search(self, request_mock):
request_mock.return_value = helpers.mock_response(
'{"hits": {"hits": []}}'
)
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
search = Search(using=client, index="test-index").filter(
"term", author="testing"
)
search.execute()
spans = self.get_finished_spans()
span = spans[0]
self.assertEqual(1, len(spans))
self.assertEqual(span.name, "Elasticsearch/<target>/_search")
self.assertIsNotNone(span.end_time)
self.assertEqual(
span.attributes,
self.search_attributes,
)
def test_dsl_search_sanitized(self, request_mock):
request_mock.return_value = helpers.mock_response(
'{"hits": {"hits": []}}'
)
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
search = Search(using=client, index="test-index").filter(
"term", author="testing"
)
search.execute()
spans = self.get_finished_spans()
span = spans[0]
self.assertEqual(1, len(spans))
self.assertEqual(span.name, "Elasticsearch/<target>/_search")
self.assertIsNotNone(span.end_time)
self.assertEqual(
span.attributes,
self.search_attributes,
)
def test_dsl_create(self, request_mock):
request_mock.side_effect = [
helpers.mock_response("{}", status_code=404),
helpers.mock_response("{}"),
]
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
Article.init(using=client)
spans = self.get_finished_spans()
assert spans
self.assertEqual(2, len(spans))
span1 = spans.by_attr(key="elasticsearch.method", value="HEAD")
span2 = spans.by_attr(key="elasticsearch.method", value="PUT")
self.assertEqual(
span1.attributes,
self.create_attributes,
)
attributes = {
DB_SYSTEM: "elasticsearch",
"elasticsearch.url": "/test-index",
"elasticsearch.method": "PUT",
}
self.assertSpanHasAttributes(span2, attributes)
self.assertEqual(
literal_eval(span2.attributes[DB_STATEMENT]),
helpers.dsl_create_statement,
)
def test_dsl_create_sanitized(self, request_mock):
request_mock.side_effect = [
helpers.mock_response("{}", status_code=404),
helpers.mock_response("{}"),
]
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
Article.init(using=client)
spans = self.get_finished_spans()
assert spans
self.assertEqual(2, len(spans))
span = spans.by_attr(key="elasticsearch.method", value="HEAD")
self.assertEqual(
span.attributes,
self.create_attributes,
)
def test_dsl_index(self, request_mock):
request_mock.return_value = helpers.mock_response(
helpers.dsl_index_result[2]
)
client = get_elasticsearch_client(hosts=["http://localhost:9200"])
article = Article(
meta={"id": 2},
title="About searching",
body="A few words here, a few words there",
)
res = article.save(using=client)
self.assertTrue(res)
spans = self.get_finished_spans()
self.assertEqual(1, len(spans))
span = spans[0]
self.assertEqual(span.name, helpers.dsl_index_span_name)
attributes = {
DB_SYSTEM: "elasticsearch",
"elasticsearch.url": helpers.dsl_index_url,
"elasticsearch.method": "PUT",
}
self.assertSpanHasAttributes(span, attributes)
self.assertEqual(
literal_eval(span.attributes[DB_STATEMENT]),
{
"body": "A few words here, a few words there",
"title": "About searching",
},
)
def test_request_hook(self, request_mock):
request_hook_method_attribute = "request_hook.method"
request_hook_url_attribute = "request_hook.url"
request_hook_kwargs_attribute = "request_hook.kwargs"
def request_hook(span, method, url, kwargs):
attributes = {
request_hook_method_attribute: method,
request_hook_url_attribute: url,
request_hook_kwargs_attribute: json.dumps(kwargs),
}
if span and span.is_recording():
span.set_attributes(attributes)
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument(request_hook=request_hook)
request_mock.return_value = helpers.mock_response(
'{"found": false, "timed_out": true, "took": 7}'
)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
index = "test-index"
doc_id = 1
kwargs = {"params": {"refresh": True, "realtime": True}}
es.get(
index=index,
id=doc_id,
**normalize_arguments(doc_type="_doc"),
**kwargs,
)
spans = self.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual(
"GET", spans[0].attributes[request_hook_method_attribute]
)
expected_url = f"/{index}/_doc/{doc_id}"
if major_version == 8:
expected_url += "?realtime=true&refresh=true"
self.assertEqual(
expected_url,
spans[0].attributes[request_hook_url_attribute],
)
if major_version == 8:
expected_kwargs = {
"body": None,
"request_timeout": "<DEFAULT>",
"max_retries": "<DEFAULT>",
"retry_on_status": "<DEFAULT>",
"retry_on_timeout": "<DEFAULT>",
"client_meta": "<DEFAULT>",
"headers": {
"accept": "application/vnd.elasticsearch+json; compatible-with=8"
},
"otel_span": None,
}
elif major_version == 7:
expected_kwargs = {
**kwargs,
"headers": {"accept": "application/json"},
}
else:
expected_kwargs = {**kwargs}
self.assertEqual(
expected_kwargs,
json.loads(spans[0].attributes[request_hook_kwargs_attribute]),
)
def test_response_hook(self, request_mock):
response_attribute_name = "db.query_result"
def response_hook(span, response):
if span and span.is_recording():
span.set_attribute(
response_attribute_name, json.dumps(response)
)
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument(response_hook=response_hook)
response_payload = {
"took": 9,
"timed_out": False,
"_shards": {
"total": 1,
"successful": 1,
"skipped": 0,
"failed": 0,
},
"hits": {
"total": {"value": 1, "relation": "eq"},
"max_score": 0.18232156,
"hits": [
{
"_index": "test-index",
"_type": "_doc",
"_id": "1",
"_score": 0.18232156,
"_source": {"name": "tester"},
}
],
},
}
request_mock.return_value = helpers.mock_response(
json.dumps(response_payload)
)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.get(
index="test-index", **normalize_arguments(doc_type="_doc"), id=1
)
spans = self.get_finished_spans()
self.assertEqual(1, len(spans))
self.assertEqual(
json.dumps(response_payload),
spans[0].attributes[response_attribute_name],
)
def test_no_op_tracer_provider(self, request_mock):
ElasticsearchInstrumentor().uninstrument()
ElasticsearchInstrumentor().instrument(
tracer_provider=trace.NoOpTracerProvider()
)
response_payload = '{"found": false, "timed_out": true, "took": 7}'
request_mock.return_value = helpers.mock_response(response_payload)
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
res = es.get(
index="test-index", **normalize_arguments(doc_type="_doc"), id=1
)
self.assertEqual(
res.get("found"), json.loads(response_payload).get("found")
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 0)
def test_body_sanitization(self, _):
self.assertEqual(
sanitize_body(sanitization_queries.interval_query),
str(sanitization_queries.interval_query_sanitized),
)
self.assertEqual(
sanitize_body(sanitization_queries.match_query),
str(sanitization_queries.match_query_sanitized),
)
self.assertEqual(
sanitize_body(sanitization_queries.filter_query),
str(sanitization_queries.filter_query_sanitized),
)
self.assertEqual(
sanitize_body(json.dumps(sanitization_queries.interval_query)),
str(sanitization_queries.interval_query_sanitized),
)
def test_bulk(self, request_mock):
request_mock.return_value = helpers.mock_response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.bulk(
body=[
{
"_op_type": "index",
"_index": "sw",
"_doc_type": "_doc",
"_id": 1,
"doc": {"name": "adam"},
},
{
"_op_type": "index",
"_index": "sw",
"_doc_type": "_doc",
"_id": 1,
"doc": {"name": "adam"},
},
]
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
# Check version and name in span's instrumentation info
self.assertEqualSpanInstrumentationScope(
span, opentelemetry.instrumentation.elasticsearch
)
@mark.skipif(
(major_version, minor_version) < (8, 13),
reason="Native OTel since elasticsearch 8.13",
)
@mock.patch.dict(
os.environ,
{"OTEL_PYTHON_INSTRUMENTATION_ELASTICSEARCH_ENABLED": "true"},
)
def test_instrumentation_is_disabled_if_native_support_enabled(
self, request_mock
):
request_mock.return_value = helpers.mock_response("{}")
es = get_elasticsearch_client(hosts=["http://localhost:9200"])
es.index(
index="sw",
id=1,
**normalize_arguments(body={"name": "adam"}, doc_type="_doc"),
)
spans_list = self.get_finished_spans()
self.assertEqual(len(spans_list), 1)
span = spans_list[0]
# Check that name in span's instrumentation info is not from this instrumentation
self.assertEqual(span.instrumentation_info.name, "elasticsearch-api")