We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticapm/instrumentation/packages/elasticsearch.py`
Content:
```
1 #  BSD 3-Clause License
2 #
3 #  Copyright (c) 2019, Elasticsearch BV
4 #  All rights reserved.
5 #
6 #  Redistribution and use in source and binary forms, with or without
7 #  modification, are permitted provided that the following conditions are met:
8 #
9 #  * Redistributions of source code must retain the above copyright notice, this
10 #    list of conditions and the following disclaimer.
11 #
12 #  * Redistributions in binary form must reproduce the above copyright notice,
13 #    this list of conditions and the following disclaimer in the documentation
14 #    and/or other materials provided with the distribution.
15 #
16 #  * Neither the name of the copyright holder nor the names of its
17 #    contributors may be used to endorse or promote products derived from
18 #    this software without specific prior written permission.
19 #
20 #  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21 #  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22 #  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23 #  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24 #  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25 #  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26 #  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30 
31 from __future__ import absolute_import
32 
33 import json
34 
35 import elasticapm
36 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
37 from elasticapm.utils import compat
38 from elasticapm.utils.logging import get_logger
39 
40 logger = get_logger("elasticapm.instrument")
41 
42 
43 API_METHOD_KEY_NAME = "__elastic_apm_api_method_name"
44 BODY_REF_NAME = "__elastic_apm_body_ref"
45 
46 
47 class ElasticSearchConnectionMixin(object):
48     query_methods = ("search", "count", "delete_by_query")
49 
50     def get_signature(self, args, kwargs):
51         args_len = len(args)
52         http_method = args[0] if args_len else kwargs.get("method")
53         http_path = args[1] if args_len > 1 else kwargs.get("url")
54 
55         return "ES %s %s" % (http_method, http_path)
56 
57     def get_context(self, instance, args, kwargs):
58         args_len = len(args)
59         params = args[2] if args_len > 2 else kwargs.get("params")
60         body = params.pop(BODY_REF_NAME, None) if params else None
61 
62         api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None
63         context = {"db": {"type": "elasticsearch"}}
64         if api_method in self.query_methods:
65             query = []
66             # using both q AND body is allowed in some API endpoints / ES versions,
67             # but not in others. We simply capture both if they are there so the
68             # user can see it.
69             if params and "q" in params:
70                 # 'q' is already encoded to a byte string at this point
71                 # we assume utf8, which is the default
72                 query.append("q=" + params["q"].decode("utf-8", errors="replace"))
73             if body and isinstance(body, dict):
74                 query.append(json.dumps(body, default=compat.text_type))
75             if query:
76                 context["db"]["statement"] = "\n\n".join(query)
77         elif api_method == "update":
78             if isinstance(body, dict) and "script" in body:
79                 # only get the `script` field from the body
80                 context["db"]["statement"] = json.dumps({"script": body["script"]})
81         context["destination"] = {
82             "address": instance.host,
83             "service": {"name": "elasticsearch", "resource": "elasticsearch", "type": "db"},
84         }
85         return context
86 
87 
88 class ElasticsearchConnectionInstrumentation(ElasticSearchConnectionMixin, AbstractInstrumentedModule):
89     name = "elasticsearch_connection"
90 
91     instrument_list = [
92         ("elasticsearch.connection.http_urllib3", "Urllib3HttpConnection.perform_request"),
93         ("elasticsearch.connection.http_requests", "RequestsHttpConnection.perform_request"),
94     ]
95 
96     def call(self, module, method, wrapped, instance, args, kwargs):
97         signature = self.get_signature(args, kwargs)
98         context = self.get_context(instance, args, kwargs)
99 
100         with elasticapm.capture_span(
101             signature,
102             span_type="db",
103             span_subtype="elasticsearch",
104             span_action="query",
105             extra=context,
106             skip_frames=2,
107             leaf=True,
108         ):
109             return wrapped(*args, **kwargs)
110 
111 
112 class ElasticsearchInstrumentation(AbstractInstrumentedModule):
113     name = "elasticsearch"
114 
115     instrument_list = [
116         ("elasticsearch.client", "Elasticsearch.delete_by_query"),
117         ("elasticsearch.client", "Elasticsearch.search"),
118         ("elasticsearch.client", "Elasticsearch.count"),
119         ("elasticsearch.client", "Elasticsearch.update"),
120     ]
121 
122     def __init__(self):
123         super(ElasticsearchInstrumentation, self).__init__()
124         try:
125             from elasticsearch import VERSION
126 
127             self.version = VERSION[0]
128         except ImportError:
129             self.version = None
130 
131     def instrument(self):
132         if self.version and not 2 <= self.version < 8:
133             logger.debug("Instrumenting version %s of Elasticsearch is not supported by Elastic APM", self.version)
134             return
135         super(ElasticsearchInstrumentation, self).instrument()
136 
137     def call(self, module, method, wrapped, instance, args, kwargs):
138         kwargs = self.inject_apm_params(method, kwargs)
139         return wrapped(*args, **kwargs)
140 
141     def inject_apm_params(self, method, kwargs):
142         params = kwargs.pop("params", {})
143 
144         # make a copy of params in case the caller reuses them for some reason
145         params = params.copy() if params is not None else {}
146 
147         method_name = method.partition(".")[-1]
148 
149         # store a reference to the non-serialized body so we can use it in the connection layer
150         body = kwargs.get("body")
151         params[BODY_REF_NAME] = body
152         params[API_METHOD_KEY_NAME] = method_name
153 
154         kwargs["params"] = params
155         return kwargs
156
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 @@
                 # '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":

{"golden_diff": "diff --git a/elasticapm/instrumentation/packages/elasticsearch.py b/elasticapm/instrumentation/packages/elasticsearch.py\n--- a/elasticapm/instrumentation/packages/elasticsearch.py\n+++ b/elasticapm/instrumentation/packages/elasticsearch.py\n@@ -58,6 +58,7 @@\n         args_len = len(args)\n         params = args[2] if args_len > 2 else kwargs.get(\"params\")\n         body = params.pop(BODY_REF_NAME, None) if params else None\n+        body_serialized = args[3] if args_len > 3 else kwargs.get(\"body\")\n \n         api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None\n         context = {\"db\": {\"type\": \"elasticsearch\"}}\n@@ -70,8 +71,16 @@\n                 # 'q' is already encoded to a byte string at this point\n                 # we assume utf8, which is the default\n                 query.append(\"q=\" + params[\"q\"].decode(\"utf-8\", errors=\"replace\"))\n-            if body and isinstance(body, dict):\n-                query.append(json.dumps(body, default=compat.text_type))\n+            if body_serialized:\n+                if isinstance(body_serialized, bytes):\n+                    query.append(body_serialized.decode(\"utf-8\", errors=\"replace\"))\n+                else:\n+                    query.append(body_serialized)\n+            elif body and isinstance(body, dict):\n+                try:\n+                    query.append(json.dumps(body, default=compat.text_type))\n+                except TypeError:\n+                    pass\n             if query:\n                 context[\"db\"][\"statement\"] = \"\\n\\n\".join(query)\n         elif api_method == \"update\":\n", "issue": "5.10.0 \"kind of\" regression for elasticsearch instrumentation due to logging of whole body\n**Describe the bug**:\r\n\r\nThe 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).\r\n\r\nThe problem seems to have happened with #940.\r\n\r\n**To Reproduce - wrong body**\r\n\r\nhttps://gist.github.com/HenrikOssipoff/a231b45ce504691cd71a5de23a771132\r\n\r\nThis gist shows a simple Starlette application with a custom serializer.\r\n\r\nUsing the code, the following document is correctly indexed:\r\n\r\n```json\r\n{\"_index\":\"test-index\",\"_type\":\"_doc\",\"_id\":\"zrmN0nYBBrB60DkP7BVZ\",\"_score\":1.0,\"_source\":{\"numbers\":2}}\r\n```\r\n\r\nHowever, the wrong body is recorded with the agent:\r\n\r\n![Sk\u00e6rmbillede 2021-01-05 kl  13 39 05](https://user-images.githubusercontent.com/4737224/103647318-69c3f400-4f5b-11eb-8841-b26b60533efc.png)\r\n\r\n**To Reproduce - crash**\r\n\r\nThe following gist shows an application that works perfectly fine with the elasticsearch client, but crashes the apm agent:\r\n\r\nhttps://gist.github.com/HenrikOssipoff/2219d059d3156ab8b11879e1191cefe4\r\n\r\nCauses exception `TypeError: keys must be str, int, float, bool or None, not NumberObj`.\r\n\r\nI'll admit this is some weird code; this is a simplified version of what we currently use.\r\nThe 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.\r\n\r\nIt works perfectly fine, but crashes the agent.\r\n\r\n**Environment (please complete the following information)**\r\n- OS: [e.g. Linux] Linux/Mac\r\n- Python version: 3.9\r\n- Framework and version [e.g. Django 2.1]: Starlette, Django (any version, probably)\r\n- APM Server version: 7.8.0\r\n- Agent version: >= 5.10.0\r\n\r\n**Additional context**\r\n\r\nI'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\r\n\r\nThe 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.\n", "before_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n\nfrom __future__ import absolute_import\n\nimport json\n\nimport elasticapm\nfrom elasticapm.instrumentation.packages.base import AbstractInstrumentedModule\nfrom elasticapm.utils import compat\nfrom elasticapm.utils.logging import get_logger\n\nlogger = get_logger(\"elasticapm.instrument\")\n\n\nAPI_METHOD_KEY_NAME = \"__elastic_apm_api_method_name\"\nBODY_REF_NAME = \"__elastic_apm_body_ref\"\n\n\nclass ElasticSearchConnectionMixin(object):\n    query_methods = (\"search\", \"count\", \"delete_by_query\")\n\n    def get_signature(self, args, kwargs):\n        args_len = len(args)\n        http_method = args[0] if args_len else kwargs.get(\"method\")\n        http_path = args[1] if args_len > 1 else kwargs.get(\"url\")\n\n        return \"ES %s %s\" % (http_method, http_path)\n\n    def get_context(self, instance, args, kwargs):\n        args_len = len(args)\n        params = args[2] if args_len > 2 else kwargs.get(\"params\")\n        body = params.pop(BODY_REF_NAME, None) if params else None\n\n        api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None\n        context = {\"db\": {\"type\": \"elasticsearch\"}}\n        if api_method in self.query_methods:\n            query = []\n            # using both q AND body is allowed in some API endpoints / ES versions,\n            # but not in others. We simply capture both if they are there so the\n            # user can see it.\n            if params and \"q\" in params:\n                # 'q' is already encoded to a byte string at this point\n                # we assume utf8, which is the default\n                query.append(\"q=\" + params[\"q\"].decode(\"utf-8\", errors=\"replace\"))\n            if body and isinstance(body, dict):\n                query.append(json.dumps(body, default=compat.text_type))\n            if query:\n                context[\"db\"][\"statement\"] = \"\\n\\n\".join(query)\n        elif api_method == \"update\":\n            if isinstance(body, dict) and \"script\" in body:\n                # only get the `script` field from the body\n                context[\"db\"][\"statement\"] = json.dumps({\"script\": body[\"script\"]})\n        context[\"destination\"] = {\n            \"address\": instance.host,\n            \"service\": {\"name\": \"elasticsearch\", \"resource\": \"elasticsearch\", \"type\": \"db\"},\n        }\n        return context\n\n\nclass ElasticsearchConnectionInstrumentation(ElasticSearchConnectionMixin, AbstractInstrumentedModule):\n    name = \"elasticsearch_connection\"\n\n    instrument_list = [\n        (\"elasticsearch.connection.http_urllib3\", \"Urllib3HttpConnection.perform_request\"),\n        (\"elasticsearch.connection.http_requests\", \"RequestsHttpConnection.perform_request\"),\n    ]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        signature = self.get_signature(args, kwargs)\n        context = self.get_context(instance, args, kwargs)\n\n        with elasticapm.capture_span(\n            signature,\n            span_type=\"db\",\n            span_subtype=\"elasticsearch\",\n            span_action=\"query\",\n            extra=context,\n            skip_frames=2,\n            leaf=True,\n        ):\n            return wrapped(*args, **kwargs)\n\n\nclass ElasticsearchInstrumentation(AbstractInstrumentedModule):\n    name = \"elasticsearch\"\n\n    instrument_list = [\n        (\"elasticsearch.client\", \"Elasticsearch.delete_by_query\"),\n        (\"elasticsearch.client\", \"Elasticsearch.search\"),\n        (\"elasticsearch.client\", \"Elasticsearch.count\"),\n        (\"elasticsearch.client\", \"Elasticsearch.update\"),\n    ]\n\n    def __init__(self):\n        super(ElasticsearchInstrumentation, self).__init__()\n        try:\n            from elasticsearch import VERSION\n\n            self.version = VERSION[0]\n        except ImportError:\n            self.version = None\n\n    def instrument(self):\n        if self.version and not 2 <= self.version < 8:\n            logger.debug(\"Instrumenting version %s of Elasticsearch is not supported by Elastic APM\", self.version)\n            return\n        super(ElasticsearchInstrumentation, self).instrument()\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        kwargs = self.inject_apm_params(method, kwargs)\n        return wrapped(*args, **kwargs)\n\n    def inject_apm_params(self, method, kwargs):\n        params = kwargs.pop(\"params\", {})\n\n        # make a copy of params in case the caller reuses them for some reason\n        params = params.copy() if params is not None else {}\n\n        method_name = method.partition(\".\")[-1]\n\n        # store a reference to the non-serialized body so we can use it in the connection layer\n        body = kwargs.get(\"body\")\n        params[BODY_REF_NAME] = body\n        params[API_METHOD_KEY_NAME] = method_name\n\n        kwargs[\"params\"] = params\n        return kwargs\n", "path": "elasticapm/instrumentation/packages/elasticsearch.py"}], "after_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n\nfrom __future__ import absolute_import\n\nimport json\n\nimport elasticapm\nfrom elasticapm.instrumentation.packages.base import AbstractInstrumentedModule\nfrom elasticapm.utils import compat\nfrom elasticapm.utils.logging import get_logger\n\nlogger = get_logger(\"elasticapm.instrument\")\n\n\nAPI_METHOD_KEY_NAME = \"__elastic_apm_api_method_name\"\nBODY_REF_NAME = \"__elastic_apm_body_ref\"\n\n\nclass ElasticSearchConnectionMixin(object):\n    query_methods = (\"search\", \"count\", \"delete_by_query\")\n\n    def get_signature(self, args, kwargs):\n        args_len = len(args)\n        http_method = args[0] if args_len else kwargs.get(\"method\")\n        http_path = args[1] if args_len > 1 else kwargs.get(\"url\")\n\n        return \"ES %s %s\" % (http_method, http_path)\n\n    def get_context(self, instance, args, kwargs):\n        args_len = len(args)\n        params = args[2] if args_len > 2 else kwargs.get(\"params\")\n        body = params.pop(BODY_REF_NAME, None) if params else None\n        body_serialized = args[3] if args_len > 3 else kwargs.get(\"body\")\n\n        api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None\n        context = {\"db\": {\"type\": \"elasticsearch\"}}\n        if api_method in self.query_methods:\n            query = []\n            # using both q AND body is allowed in some API endpoints / ES versions,\n            # but not in others. We simply capture both if they are there so the\n            # user can see it.\n            if params and \"q\" in params:\n                # 'q' is already encoded to a byte string at this point\n                # we assume utf8, which is the default\n                query.append(\"q=\" + params[\"q\"].decode(\"utf-8\", errors=\"replace\"))\n            if body_serialized:\n                if isinstance(body_serialized, bytes):\n                    query.append(body_serialized.decode(\"utf-8\", errors=\"replace\"))\n                else:\n                    query.append(body_serialized)\n            elif body and isinstance(body, dict):\n                try:\n                    query.append(json.dumps(body, default=compat.text_type))\n                except TypeError:\n                    pass\n            if query:\n                context[\"db\"][\"statement\"] = \"\\n\\n\".join(query)\n        elif api_method == \"update\":\n            if isinstance(body, dict) and \"script\" in body:\n                # only get the `script` field from the body\n                context[\"db\"][\"statement\"] = json.dumps({\"script\": body[\"script\"]})\n        context[\"destination\"] = {\n            \"address\": instance.host,\n            \"service\": {\"name\": \"elasticsearch\", \"resource\": \"elasticsearch\", \"type\": \"db\"},\n        }\n        return context\n\n\nclass ElasticsearchConnectionInstrumentation(ElasticSearchConnectionMixin, AbstractInstrumentedModule):\n    name = \"elasticsearch_connection\"\n\n    instrument_list = [\n        (\"elasticsearch.connection.http_urllib3\", \"Urllib3HttpConnection.perform_request\"),\n        (\"elasticsearch.connection.http_requests\", \"RequestsHttpConnection.perform_request\"),\n    ]\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        signature = self.get_signature(args, kwargs)\n        context = self.get_context(instance, args, kwargs)\n\n        with elasticapm.capture_span(\n            signature,\n            span_type=\"db\",\n            span_subtype=\"elasticsearch\",\n            span_action=\"query\",\n            extra=context,\n            skip_frames=2,\n            leaf=True,\n        ):\n            return wrapped(*args, **kwargs)\n\n\nclass ElasticsearchInstrumentation(AbstractInstrumentedModule):\n    name = \"elasticsearch\"\n\n    instrument_list = [\n        (\"elasticsearch.client\", \"Elasticsearch.delete_by_query\"),\n        (\"elasticsearch.client\", \"Elasticsearch.search\"),\n        (\"elasticsearch.client\", \"Elasticsearch.count\"),\n        (\"elasticsearch.client\", \"Elasticsearch.update\"),\n    ]\n\n    def __init__(self):\n        super(ElasticsearchInstrumentation, self).__init__()\n        try:\n            from elasticsearch import VERSION\n\n            self.version = VERSION[0]\n        except ImportError:\n            self.version = None\n\n    def instrument(self):\n        if self.version and not 2 <= self.version < 8:\n            logger.debug(\"Instrumenting version %s of Elasticsearch is not supported by Elastic APM\", self.version)\n            return\n        super(ElasticsearchInstrumentation, self).instrument()\n\n    def call(self, module, method, wrapped, instance, args, kwargs):\n        kwargs = self.inject_apm_params(method, kwargs)\n        return wrapped(*args, **kwargs)\n\n    def inject_apm_params(self, method, kwargs):\n        params = kwargs.pop(\"params\", {})\n\n        # make a copy of params in case the caller reuses them for some reason\n        params = params.copy() if params is not None else {}\n\n        method_name = method.partition(\".\")[-1]\n\n        # store a reference to the non-serialized body so we can use it in the connection layer\n        body = kwargs.get(\"body\")\n        params[BODY_REF_NAME] = body\n        params[API_METHOD_KEY_NAME] = method_name\n\n        kwargs[\"params\"] = params\n        return kwargs\n", "path": "elasticapm/instrumentation/packages/elasticsearch.py"}]}