Fix values sanitizer handling queries with keys that contain dots

Author: imamatory

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)

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": {}}},
+}

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("{}")