Repository: elastic/apm-agent-python

Issue:
APM adding query params causing 400 bad request error during ES query
**Bug description**

I have an API that queries ElasticSearch and that I'm trying to monitor thanks to APM. Somehow I have the exact same problem described in this (quite old) discussion:
https://discuss.elastic.co/t/apm-adding-query-params-causing-400-bad-request-error-during-es-query/186749
```
elasticsearch.exceptions.RequestError: RequestError(400, 'illegal_argument_exception', 'request [/my_index/_search] contains unrecognized parameters: [__elastic_apm_api_method_name], [__elastic_apm_body_ref]')
```
The workaround described in this discussion (putting `SKIP_INSTRUMENT_ELASTICSEARCH=true` in my environment variables) does work, but I would still like to instrument my ElasticSearch calls.

**Context**
I have the latest version of everything:
* a docker container with Elastic 7.4.2
* my Python code runs APM 5.3.2. 

To be noted that I use the Async Python Client for Elastic, whose latest version is only 6.2.0, maybe this could lead to some inconsistencies. 



Discussion:
Hi @Thematrixme 

It does indeed look like this is caused by the use of the async client, or more precisely, our lack of support for it.

The good news is that we recently merged initial support for asyncio, so it should be possible to add instrumentation for the async client in a more or less straight forward way. I'll give it a try and report back my findings :) 
@Thematrixme quick follow up, what Python version do you run? For our asyncio support, we currently only support 3.7+
@beniwohli thanks for both your answers. I use Python 3.7 for this project, so I'm good. I'm looking forward to your solution!

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/asyncio/elasticsearch.py b/elasticapm/instrumentation/packages/asyncio/elasticsearch.py
new file mode 100644
--- /dev/null
+++ b/elasticapm/instrumentation/packages/asyncio/elasticsearch.py
@@ -0,0 +1,54 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import elasticapm
+from elasticapm.instrumentation.packages.asyncio.base import AsyncAbstractInstrumentedModule
+from elasticapm.instrumentation.packages.elasticsearch import ElasticSearchConnectionMixin
+
+
+class ElasticSearchAsyncConnection(ElasticSearchConnectionMixin, AsyncAbstractInstrumentedModule):
+    name = "elasticsearch_connection"
+
+    instrument_list = [("elasticsearch_async.connection", "AIOHttpConnection.perform_request")]
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
+        signature = self.get_signature(args, kwargs)
+        context = self.get_context(args, kwargs)
+
+        async with elasticapm.async_capture_span(
+            signature,
+            span_type="db",
+            span_subtype="elasticsearch",
+            span_action="query",
+            extra=context,
+            skip_frames=2,
+            leaf=True,
+        ):
+            return await wrapped(*args, **kwargs)
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
@@ -44,26 +44,22 @@
 BODY_REF_NAME = "__elastic_apm_body_ref"
 
 
-class ElasticsearchConnectionInstrumentation(AbstractInstrumentedModule):
-    name = "elasticsearch_connection"
-
-    instrument_list = [
-        ("elasticsearch.connection.http_urllib3", "Urllib3HttpConnection.perform_request"),
-        ("elasticsearch.connection.http_requests", "RequestsHttpConnection.perform_request"),
-    ]
-
+class ElasticSearchConnectionMixin(object):
     query_methods = ("Elasticsearch.search", "Elasticsearch.count", "Elasticsearch.delete_by_query")
 
-    def call(self, module, method, wrapped, instance, args, kwargs):
+    def get_signature(self, args, kwargs):
         args_len = len(args)
         http_method = args[0] if args_len else kwargs.get("method")
         http_path = args[1] if args_len > 1 else kwargs.get("url")
+
+        return "ES %s %s" % (http_method, http_path)
+
+    def get_context(self, 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
 
         api_method = params.pop(API_METHOD_KEY_NAME, None) if params else None
-
-        signature = "ES %s %s" % (http_method, http_path)
         context = {"db": {"type": "elasticsearch"}}
         if api_method in self.query_methods:
             query = []
@@ -83,6 +79,21 @@ def call(self, module, method, wrapped, instance, args, kwargs):
                 # only get the `script` field from the body
                 context["db"]["statement"] = json.dumps({"script": body["script"]})
         # TODO: add instance.base_url to context once we agreed on a format
+        return context
+
+
+class ElasticsearchConnectionInstrumentation(ElasticSearchConnectionMixin, AbstractInstrumentedModule):
+    name = "elasticsearch_connection"
+
+    instrument_list = [
+        ("elasticsearch.connection.http_urllib3", "Urllib3HttpConnection.perform_request"),
+        ("elasticsearch.connection.http_requests", "RequestsHttpConnection.perform_request"),
+    ]
+
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        signature = self.get_signature(args, kwargs)
+        context = self.get_context(args, kwargs)
+
         with elasticapm.capture_span(
             signature,
             span_type="db",
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/elasticapm/instrumentation/register.py
@@ -65,6 +65,7 @@
         [
             "elasticapm.instrumentation.packages.asyncio.sleep.AsyncIOSleepInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiohttp_client.AioHttpClientInstrumentation",
+            "elasticapm.instrumentation.packages.asyncio.elasticsearch.ElasticSearchAsyncConnection",
             "elasticapm.instrumentation.packages.asyncio.aiopg.AioPGInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoRequestExecuteInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoHandleRequestExceptionInstrumentation",


Test patch:
diff --git a/tests/instrumentation/asyncio/elasticsearch_async_client_tests.py b/tests/instrumentation/asyncio/elasticsearch_async_client_tests.py
new file mode 100644
--- /dev/null
+++ b/tests/instrumentation/asyncio/elasticsearch_async_client_tests.py
@@ -0,0 +1,190 @@
+#  BSD 3-Clause License
+#
+#  Copyright (c) 2019, Elasticsearch BV
+#  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#
+#  * Redistributions of source code must retain the above copyright notice, this
+#    list of conditions and the following disclaimer.
+#
+#  * Redistributions in binary form must reproduce the above copyright notice,
+#    this list of conditions and the following disclaimer in the documentation
+#    and/or other materials provided with the distribution.
+#
+#  * Neither the name of the copyright holder nor the names of its
+#    contributors may be used to endorse or promote products derived from
+#    this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+#  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import pytest  # isort:skip
+
+pytest.importorskip("elasticsearch_async")  # isort:skip
+
+import os
+
+from elasticsearch import VERSION as ES_VERSION
+
+from elasticapm.conf.constants import TRANSACTION
+from elasticsearch_async import AsyncElasticsearch
+
+pytestmark = [pytest.mark.elasticsearch, pytest.mark.asyncio, pytest.mark.integrationtest]
+
+if "ES_URL" not in os.environ:
+    pytestmark.append(pytest.mark.skip("Skipping elasticsearch test, no ES_URL environment variable"))
+
+
+document_type = "_doc" if ES_VERSION[0] >= 6 else "doc"
+
+
+@pytest.fixture
+async def elasticsearch_async(request):
+    """Elasticsearch client fixture."""
+    client = AsyncElasticsearch(hosts=os.environ["ES_URL"])
+    try:
+        yield client
+    finally:
+        await client.indices.delete(index="*")
+        await client.transport.close()
+
+
+async def test_ping(instrument, elasticapm_client, elasticsearch_async):
+    elasticapm_client.begin_transaction("test")
+    result = await elasticsearch_async.ping()
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["name"] == "ES HEAD /"
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["sync"] is False
+
+
+async def test_info(instrument, elasticapm_client, elasticsearch_async):
+    elasticapm_client.begin_transaction("test")
+    result = await elasticsearch_async.info()
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["name"] == "ES GET /"
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["sync"] is False
+
+
+async def test_create(instrument, elasticapm_client, elasticsearch_async):
+    elasticapm_client.begin_transaction("test")
+    if ES_VERSION[0] < 5:
+        r1 = await elasticsearch_async.create("tweets", document_type, {"user": "kimchy", "text": "hola"}, 1)
+    elif ES_VERSION[0] < 7:
+        r1 = await elasticsearch_async.create("tweets", document_type, 1, body={"user": "kimchy", "text": "hola"})
+    else:
+        r1 = await elasticsearch_async.create("tweets", 1, body={"user": "kimchy", "text": "hola"})
+    r2 = await elasticsearch_async.create(
+        index="tweets", doc_type=document_type, id=2, body={"user": "kimchy", "text": "hola"}, refresh=True
+    )
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 2
+
+    for i, span in enumerate(spans):
+        if ES_VERSION[0] >= 5:
+            assert span["name"] == "ES PUT /tweets/%s/%d/_create" % (document_type, i + 1)
+        else:
+            assert span["name"] == "ES PUT /tweets/%s/%d" % (document_type, i + 1)
+        assert span["type"] == "db"
+        assert span["subtype"] == "elasticsearch"
+        assert span["action"] == "query"
+        assert span["context"]["db"]["type"] == "elasticsearch"
+        assert "statement" not in span["context"]["db"]
+
+
+async def test_search_body(instrument, elasticapm_client, elasticsearch_async):
+    await elasticsearch_async.create(
+        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
+    )
+    elasticapm_client.begin_transaction("test")
+    search_query = {"query": {"term": {"user": "kimchy"}}}
+    result = await elasticsearch_async.search(body=search_query, params=None)
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    assert result["hits"]["hits"][0]["_source"] == {"user": "kimchy", "text": "hola"}
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["name"] in ("ES GET /_search", "ES GET /_all/_search")
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}'
+    assert span["sync"] is False
+
+
+async def test_count_body(instrument, elasticapm_client, elasticsearch_async):
+    await elasticsearch_async.create(
+        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
+    )
+    elasticapm_client.begin_transaction("test")
+    search_query = {"query": {"term": {"user": "kimchy"}}}
+    result = await elasticsearch_async.count(body=search_query)
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    assert result["count"] == 1
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+    assert span["name"] in ("ES GET /_count", "ES POST /_count", "ES GET /_all/_count")
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}'
+    assert span["sync"] is False
+
+
+@pytest.mark.skipif(ES_VERSION[0] < 5, reason="unsupported method")
+async def test_delete_by_query_body(instrument, elasticapm_client, elasticsearch_async):
+    await elasticsearch_async.create(
+        index="tweets", doc_type=document_type, id=1, body={"user": "kimchy", "text": "hola"}, refresh=True
+    )
+    elasticapm_client.begin_transaction("test")
+    result = await elasticsearch_async.delete_by_query(index="tweets", body={"query": {"term": {"user": "kimchy"}}})
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    spans = elasticapm_client.spans_for_transaction(transaction)
+
+    span = spans[0]
+    assert span["name"] == "ES POST /tweets/_delete_by_query"
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["db"]["statement"] == '{"term": {"user": "kimchy"}}'
+    assert span["sync"] is False
diff --git a/tests/requirements/requirements-elasticsearch-6.txt b/tests/requirements/requirements-elasticsearch-6.txt
--- a/tests/requirements/requirements-elasticsearch-6.txt
+++ b/tests/requirements/requirements-elasticsearch-6.txt
@@ -1,2 +1,3 @@
 elasticsearch>=6.0,<7.0
+elasticsearch-async ; python_version >= '3.7'
 -r requirements-base.txt
diff --git a/tests/requirements/requirements-elasticsearch-7.txt b/tests/requirements/requirements-elasticsearch-7.txt
--- a/tests/requirements/requirements-elasticsearch-7.txt
+++ b/tests/requirements/requirements-elasticsearch-7.txt
@@ -1,2 +1,3 @@
 elasticsearch>=7.0,<8.0
+elasticsearch-async ; python_version >= '3.7'
 -r requirements-base.txt