diff --git a/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/utils.py b/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/utils.py index 97f2bc3b87..819d30cf94 100644 --- a/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/utils.py +++ b/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/utils.py @@ -25,40 +25,22 @@ sanitized_value = "?" -# pylint: disable=C0103 -def _flatten_dict(d, parent_key=""): - items = [] - for k, v in d.items(): - new_key = parent_key + "." + k if parent_key else k - # recursive call _flatten_dict for a non-empty dict value - if isinstance(v, dict) and v: - items.extend(_flatten_dict(v, new_key).items()) - else: - items.append((new_key, v)) - return dict(items) - -def _unflatten_dict(d): - res = {} - for k, v in d.items(): - keys = k.split(".") - d = res - for key in keys[:-1]: - if key not in d: - d[key] = {} - d = d[key] - d[keys[-1]] = v - return res +def sanitize_dict(d): + sanitized_copy = {} + for key, value in d.items(): + if isinstance(value, dict): + sanitized_copy[key] = sanitize_dict(value) + elif key in sanitized_keys: + sanitized_copy[key] = sanitized_value + else: + sanitized_copy[key] = value + return sanitized_copy def sanitize_body(body) -> str: if isinstance(body, str): body = json.loads(body) - flatten_body = _flatten_dict(body) - - for key in flatten_body: - if key.endswith(sanitized_keys): - flatten_body[key] = sanitized_value - - return str(_unflatten_dict(flatten_body)) + sanitized_body = sanitize_dict(body) + return str(sanitized_body) diff --git a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/sanitization_queries.py b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/sanitization_queries.py index 234e24433e..79841117db 100644 --- a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/sanitization_queries.py +++ b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/sanitization_queries.py @@ -44,6 +44,16 @@ } } +_field_opts = {} +highlight_query = { + "query": { + "match": {"content": "search"}, + }, + "highlight": { + "fields": {"content": _field_opts, "content.morph": _field_opts}, + }, +} + interval_query_sanitized = { "query": { "intervals": { @@ -63,3 +73,8 @@ } } } + +highlight_query_sanitized = { + "query": {"match": {"content": "search"}}, + "highlight": {"fields": {"content": {}, "content.morph": {}}}, +} diff --git a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py index 8a707c8c2a..0f9b253d66 100644 --- a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py +++ b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py @@ -579,6 +579,10 @@ def test_body_sanitization(self, _): sanitize_body(json.dumps(sanitization_queries.interval_query)), str(sanitization_queries.interval_query_sanitized), ) + self.assertEqual( + sanitize_body(sanitization_queries.highlight_query), + str(sanitization_queries.highlight_query_sanitized), + ) def test_bulk(self, request_mock): request_mock.return_value = helpers.mock_response("{}")