Fix for json-incompatible elasticsearch bodies (#998)

* Use the elasticsearch serializer for better compat

* Use the `body` provided with perform_request

* Check for bytes body

* CHANGELOG

Author: basepi

CHANGELOG.asciidoc

@@ -30,6 +30,22 @@ endif::[]
 //===== Bug fixes
 //
 
+=== Unreleased
+
+// Unreleased changes go here
+// When the next release happens, nest these changes under the "Python Agent version 5.x" heading
+[float]
+===== Features
+
+
+[float]
+===== Bug fixes
+
+* Fix for custom serializers in elasticsearch-py {pull}998[#998]
+* Fix large query truncation in psycopg2 {pull}994[#994]
+* Fix memory metrics reporting when `memory.usage_in_bytes` is unavailable {pull}987[#987]
+
+
 [[release-notes-5.x]]
 === Python Agent version 5.x
 
@@ -38,14 +54,14 @@ endif::[]
 [float]
 ===== Features
 
- * Implement "sample_rate" property for transactions and spans, and propagate through tracestate {pull}891[#891]
- * Add support for callbacks on config changes {pull}912[#912]
- * Override `sys.excepthook` to catch all exceptions {pull}943[#943]
- * Implement `log_level` config (supports central config) {pull}946[#946]
- * Implement `transaction_ignore_urls` config (supports central config) {pull}923[#923]
- * Add public API to retrieve trace parent header {pull}956[#956]
- * Added support for cgroup memory metrics {pull}846[#846]
- 
+* Implement "sample_rate" property for transactions and spans, and propagate through tracestate {pull}891[#891]
+* Add support for callbacks on config changes {pull}912[#912]
+* Override `sys.excepthook` to catch all exceptions {pull}943[#943]
+* Implement `log_level` config (supports central config) {pull}946[#946]
+* Implement `transaction_ignore_urls` config (supports central config) {pull}923[#923]
+* Add public API to retrieve trace parent header {pull}956[#956]
+* Added support for cgroup memory metrics {pull}846[#846]
+
 
 [float]
 ===== Bug fixes
@@ -62,9 +78,9 @@ endif::[]
 [float]
 ===== Features
 
- * Add instrumentation support for https://github.com/encode/httpx[`httpx`] and https://github.com/encode/httpcore[`httpcore`] {pull}898[#898]
- * Implement "outcome" property for transactions and spans {pull}899[#899]
- * Add support for `asyncpg` {pull}889[#889]
+* Add instrumentation support for https://github.com/encode/httpx[`httpx`] and https://github.com/encode/httpcore[`httpcore`] {pull}898[#898]
+* Implement "outcome" property for transactions and spans {pull}899[#899]
+* Add support for `asyncpg` {pull}889[#889]
 
 [float]
 ===== Bug fixes

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":

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

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}}}}')