Repository: elastic/apm-agent-python

Issue:
5.10.0 "kind of" regression for elasticsearch instrumentation due to logging of whole body
**Describe the bug**:

The elasticsearch library makes it possible to define your own custom serializer as part of the library. However, the Python agent doesn't take this into account, which can lead to wrong request bodies being collected, and it is possible to downright crash the agent (which unfortunately was what happened to us).

The problem seems to have happened with #940.

**To Reproduce - wrong body**

https://gist.github.com/HenrikOssipoff/a231b45ce504691cd71a5de23a771132

This gist shows a simple Starlette application with a custom serializer.

Using the code, the following document is correctly indexed:

```json
{"_index":"test-index","_type":"_doc","_id":"zrmN0nYBBrB60DkP7BVZ","_score":1.0,"_source":{"numbers":2}}
```

However, the wrong body is recorded with the agent:

![Skærmbillede 2021-01-05 kl  13 39 05](https://user-images.githubusercontent.com/4737224/103647318-69c3f400-4f5b-11eb-8841-b26b60533efc.png)

**To Reproduce - crash**

The following gist shows an application that works perfectly fine with the elasticsearch client, but crashes the apm agent:

https://gist.github.com/HenrikOssipoff/2219d059d3156ab8b11879e1191cefe4

Causes exception `TypeError: keys must be str, int, float, bool or None, not NumberObj`.

I'll admit this is some weird code; this is a simplified version of what we currently use.
The version in our app uses Django's translation module as keys for certain parts of the query, which corresponds to `django.utils.functional.Promise` objects. We're basically looping over these to force-evaluate them, as they're lazy objects.

It works perfectly fine, but crashes the agent.

**Environment (please complete the following information)**
- OS: [e.g. Linux] Linux/Mac
- Python version: 3.9
- Framework and version [e.g. Django 2.1]: Starlette, Django (any version, probably)
- APM Server version: 7.8.0
- Agent version: >= 5.10.0

**Additional context**

I'm not actually sure how best to solve this. The elasticsearch library already supports a wide array of custom serialization, even without using your own custom one: https://github.com/elastic/elasticsearch-py/blob/master/elasticsearch/serializer.py

The best case (from a user perspective), would be for the agent to somehow find the correct serializer to use, when dealing with the body. I'm not sure if this is even possible.


Discussion:
If anyone else is needing a fix for this, the only fix we've found is to simply downgrade the python agent to <5.10.0 for now.

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/elasticsearch.py b/elasticapm/instrumentation/packages/elasticsearch.py
--- a/elasticapm/instrumentation/packages/elasticsearch.py
+++ b/elasticapm/instrumentation/packages/elasticsearch.py
@@ -58,6 +58,7 @@ def get_context(self, instance, args, kwargs):
         args_len = len(args)
         params = args[2] if args_len > 2 else kwargs.get("params")
         body = params.pop(BODY_REF_NAME, None) if params else None
+        body_serialized = args[3] if args_len > 3 else kwargs.get("body")
 
         api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None
         context = {"db": {"type": "elasticsearch"}}
@@ -70,8 +71,16 @@ def get_context(self, instance, args, kwargs):
                 # 'q' is already encoded to a byte string at this point
                 # we assume utf8, which is the default
                 query.append("q=" + params["q"].decode("utf-8", errors="replace"))
-            if body and isinstance(body, dict):
-                query.append(json.dumps(body, default=compat.text_type))
+            if body_serialized:
+                if isinstance(body_serialized, bytes):
+                    query.append(body_serialized.decode("utf-8", errors="replace"))
+                else:
+                    query.append(body_serialized)
+            elif body and isinstance(body, dict):
+                try:
+                    query.append(json.dumps(body, default=compat.text_type))
+                except TypeError:
+                    pass
             if query:
                 context["db"]["statement"] = "\n\n".join(query)
         elif api_method == "update":


Test patch:
diff --git a/tests/instrumentation/asyncio_tests/async_elasticsearch_client_tests.py b/tests/instrumentation/asyncio_tests/async_elasticsearch_client_tests.py
--- a/tests/instrumentation/asyncio_tests/async_elasticsearch_client_tests.py
+++ b/tests/instrumentation/asyncio_tests/async_elasticsearch_client_tests.py
@@ -32,6 +32,7 @@
 
 pytest.importorskip("elasticsearch._async")  # isort:skip
 
+import json
 import os
 
 from elasticsearch import VERSION as ES_VERSION
@@ -145,9 +146,10 @@ async def test_search_body(instrument, elasticapm_client, async_elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert (
-        span["context"]["db"]["statement"] == '{"sort": ["userid"], "query": {"term": {"user": "kimchy"}}}'
-        or span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads(
+        '{"sort": ["userid"], "query": {"term": {"user": "kimchy"}}}'
+    ) or json.loads(span["context"]["db"]["statement"]) == json.loads(
+        '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
     )
     assert span["sync"] is False
 
@@ -171,7 +173,7 @@ async def test_count_body(instrument, elasticapm_client, async_elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query": {"term": {"user": "kimchy"}}}')
     assert span["sync"] is False
 
 
@@ -192,5 +194,5 @@ async def test_delete_by_query_body(instrument, elasticapm_client, async_elastic
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query":{"term":{"user":"kimchy"}}}')
     assert span["sync"] is False
diff --git a/tests/instrumentation/elasticsearch_tests.py b/tests/instrumentation/elasticsearch_tests.py
--- a/tests/instrumentation/elasticsearch_tests.py
+++ b/tests/instrumentation/elasticsearch_tests.py
@@ -32,12 +32,15 @@
 
 pytest.importorskip("elasticsearch")  # isort:skip
 
+import json
 import os
 
 from elasticsearch import VERSION as ES_VERSION
 from elasticsearch import Elasticsearch
+from elasticsearch.serializer import JSONSerializer
 
 from elasticapm.conf.constants import TRANSACTION
+from elasticapm.utils import compat
 
 pytestmark = [pytest.mark.elasticsearch]
 
@@ -48,10 +51,39 @@
 document_type = "_doc" if ES_VERSION[0] >= 6 else "doc"
 
 
+class NumberObj:
+    def __init__(self, value):
+        self.value = value
+
+
+class SpecialEncoder(JSONSerializer):
+    def default(self, obj):
+        if isinstance(obj, NumberObj):
+            return obj.value
+        return JSONSerializer.default(self, obj)
+
+    def force_key_encoding(self, obj):
+        if isinstance(obj, dict):
+
+            def yield_key_value(d):
+                for key, value in compat.iteritems(d):
+                    try:
+                        yield self.default(key), self.force_key_encoding(value)
+                    except TypeError:
+                        yield key, self.force_key_encoding(value)
+
+            return dict(yield_key_value(obj))
+        else:
+            return obj
+
+    def dumps(self, obj):
+        return super(SpecialEncoder, self).dumps(self.force_key_encoding(obj))
+
+
 @pytest.fixture
 def elasticsearch(request):
     """Elasticsearch client fixture."""
-    client = Elasticsearch(hosts=os.environ["ES_URL"])
+    client = Elasticsearch(hosts=os.environ["ES_URL"], serializer=SpecialEncoder())
     try:
         yield client
     finally:
@@ -338,9 +370,10 @@ def test_search_body(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert (
-        span["context"]["db"]["statement"] == '{"sort": ["userid"], "query": {"term": {"user": "kimchy"}}}'
-        or span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads(
+        '{"sort": ["userid"], "query": {"term": {"user": "kimchy"}}}'
+    ) or json.loads(span["context"]["db"]["statement"]) == json.loads(
+        '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
     )
 
 
@@ -393,7 +426,7 @@ def test_search_both(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert span["context"]["db"]["statement"] == 'q=text:hola\n\n{"query": {"term": {"user": "kimchy"}}}'
+    assert span["context"]["db"]["statement"].startswith('q=text:hola\n\n{"query":')
 
 
 @pytest.mark.integrationtest
@@ -419,7 +452,7 @@ def test_count_body(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query": {"term": {"user": "kimchy"}}}')
 
 
 @pytest.mark.integrationtest
@@ -486,7 +519,7 @@ def test_delete_by_query_body(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
-    assert span["context"]["db"]["statement"] == '{"query": {"term": {"user": "kimchy"}}}'
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query":{"term":{"user":"kimchy"}}}')
 
 
 @pytest.mark.integrationtest
@@ -528,3 +561,20 @@ def test_multiple_indexes_doctypes(instrument, elasticapm_client, elasticsearch)
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
+
+
+@pytest.mark.integrationtest
+def test_custom_serializer(instrument, elasticapm_client, elasticsearch):
+    if ES_VERSION[0] < 7:
+        elasticsearch.index("test-index", document_type, {"2": 1})
+    else:
+        elasticsearch.index(index="test-index", body={"2": 1})
+    elasticapm_client.begin_transaction("test")
+    search_query = {"query": {"term": {NumberObj(2): {"value": 1}}}}
+    result = elasticsearch.search(index="test-index", body=search_query, params=None)
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    span = spans[0]
+    assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query":{"term":{"2":{"value":1}}}}')