Repository: elastic/apm-agent-python

Issue:
Tracking enhacements for `Elasticsearch.search`
Currently in https://github.com/elastic/apm-agent-python/blob/master/elasticapm/instrumentation/packages/elasticsearch.py#L47 we are discarding everything but the `query` in the body which can cover some very important (from a performance point of view) aspects of the operation for example pagination - large `from`/`size` can have severe impact - and aggregations - if present can severely impact everything.

It would also be great to track the statistics on the response:

`took`, `hits.total`, `_shards.total/failed/successful` information

which can expose badly structured code/queries or performance issues in elasticsearch

Thank you!


Discussion:
@HonzaKral (or @sethmlarson) can you clarify why including `aggs` can be problematic for collecting the query? I recently "fixed" this in #940 because a user was complaining about the lack of `aggs` which were the important part of the query for his debugging.

For the additional information we could collect, I'm not sure where we would put it. We currently don't populate any custom context ourselves, we leave that for the user. We also don't collect service-specific data like this for anything else we instrument.

Unless @beniwohli has any bright ideas I think I'm going to probably close this issue (once we resolve the above question about `aggs` and other query data).
I don't see how including the whole body could be problematic, quite the opposite. The `_search` API is used for many other things besides just running a `query` (aggregations, suggestions, sorting, script fields, pagination, ... ...) all of which can have **significant** impact on performance and behavior and are therefore important to track.

> We also don't collect service-specific data like this for anything else we instrument.

We do provide information for `rows_affected` for SQL for example, this would be similar and would belong in the same place. 

Ah, I misread your original issue, I thought you were saying that we purposely only include the query because including the whole payload could be problematic. Glad my fix was a good one!

And thank you for pointing out `rows_affected` on the span context. I hadn't actually internalized that spans have their own context, and it doesn't have a standalone spec in the apm server repo. Rather, span context is here: https://github.com/elastic/apm-server/blob/master/docs/spec/spans/span.json

I think `hits.total` could probably go under the existing `rows_affected`, and we could easily add something for `took`. I'm trying to figure out  what field we might use for the `_shards.total/failed/successful` information. Ideally it would be a field other db providers could use for similar context. I'll have to think on this (and we'll need a PR against that spec before we can add these new context pieces).

Resolving patch:
diff --git a/elasticapm/instrumentation/packages/asyncio/elasticsearch.py b/elasticapm/instrumentation/packages/asyncio/elasticsearch.py
--- a/elasticapm/instrumentation/packages/asyncio/elasticsearch.py
+++ b/elasticapm/instrumentation/packages/asyncio/elasticsearch.py
@@ -30,10 +30,14 @@
 
 import elasticapm
 from elasticapm.instrumentation.packages.asyncio.base import AsyncAbstractInstrumentedModule
-from elasticapm.instrumentation.packages.elasticsearch import ElasticSearchConnectionMixin
+from elasticapm.instrumentation.packages.elasticsearch import (
+    ElasticsearchConnectionInstrumentation,
+    ElasticsearchTransportInstrumentation,
+)
+from elasticapm.traces import execution_context
 
 
-class ElasticSearchAsyncConnection(ElasticSearchConnectionMixin, AsyncAbstractInstrumentedModule):
+class ElasticSearchAsyncConnection(ElasticsearchConnectionInstrumentation, AsyncAbstractInstrumentedModule):
     name = "elasticsearch_connection"
 
     instrument_list = [
@@ -42,16 +46,39 @@ class ElasticSearchAsyncConnection(ElasticSearchConnectionMixin, AsyncAbstractIn
     ]
 
     async def call(self, module, method, wrapped, instance, args, kwargs):
-        signature = self.get_signature(args, kwargs)
-        context = self.get_context(instance, args, kwargs)
+        span = execution_context.get_span()
 
+        self._update_context_by_request_data(span.context, instance, args, kwargs)
+
+        status_code, headers, raw_data = await wrapped(*args, **kwargs)
+
+        span.context["http"] = {"status_code": status_code}
+
+        return status_code, headers, raw_data
+
+
+class ElasticsearchAsyncTransportInstrumentation(
+    ElasticsearchTransportInstrumentation, AsyncAbstractInstrumentedModule
+):
+    name = "elasticsearch_connection"
+
+    instrument_list = [
+        ("elasticsearch._async.transport", "AsyncTransport.perform_request"),
+    ]
+
+    async def call(self, module, method, wrapped, instance, args, kwargs):
         async with elasticapm.async_capture_span(
-            signature,
+            self._get_signature(args, kwargs),
             span_type="db",
             span_subtype="elasticsearch",
             span_action="query",
-            extra=context,
+            extra={},
             skip_frames=2,
             leaf=True,
-        ):
-            return await wrapped(*args, **kwargs)
+        ) as span:
+            result_data = await wrapped(*args, **kwargs)
+
+            if isinstance(result_data, dict) and "hits" in result_data:
+                span.context["db"]["rows_affected"] = result_data["hits"]["total"]["value"]
+
+            return result_data
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
@@ -34,6 +34,7 @@
 
 import elasticapm
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
+from elasticapm.traces import execution_context
 from elasticapm.utils.logging import get_logger
 
 logger = get_logger("elasticapm.instrument")
@@ -41,17 +42,26 @@
 should_capture_body_re = re.compile("/(_search|_msearch|_count|_async_search|_sql|_eql)(/|$)")
 
 
-class ElasticSearchConnectionMixin(object):
-    query_methods = ("search", "count", "delete_by_query")
+class ElasticsearchConnectionInstrumentation(AbstractInstrumentedModule):
+    name = "elasticsearch_connection"
 
-    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")
+    instrument_list = [
+        ("elasticsearch.connection.http_urllib3", "Urllib3HttpConnection.perform_request"),
+        ("elasticsearch.connection.http_requests", "RequestsHttpConnection.perform_request"),
+    ]
 
-        return "ES %s %s" % (http_method, http_path)
+    def call(self, module, method, wrapped, instance, args, kwargs):
+        span = execution_context.get_span()
+
+        self._update_context_by_request_data(span.context, instance, args, kwargs)
+
+        status_code, headers, raw_data = wrapped(*args, **kwargs)
+
+        span.context["http"] = {"status_code": status_code}
 
-    def get_context(self, instance, args, kwargs):
+        return status_code, headers, raw_data
+
+    def _update_context_by_request_data(self, context, instance, args, kwargs):
         args_len = len(args)
         url = args[1] if args_len > 1 else kwargs.get("url")
         params = args[2] if args_len > 2 else kwargs.get("params")
@@ -59,7 +69,7 @@ def get_context(self, instance, args, kwargs):
 
         should_capture_body = bool(should_capture_body_re.search(url))
 
-        context = {"db": {"type": "elasticsearch"}}
+        context["db"] = {"type": "elasticsearch"}
         if should_capture_body:
             query = []
             # using both q AND body is allowed in some API endpoints / ES versions,
@@ -76,32 +86,42 @@ def get_context(self, instance, args, kwargs):
                     query.append(body_serialized)
             if query:
                 context["db"]["statement"] = "\n\n".join(query)
+
         context["destination"] = {
             "address": instance.host,
             "service": {"name": "elasticsearch", "resource": "elasticsearch", "type": "db"},
         }
-        return context
 
 
-class ElasticsearchConnectionInstrumentation(ElasticSearchConnectionMixin, AbstractInstrumentedModule):
+class ElasticsearchTransportInstrumentation(AbstractInstrumentedModule):
     name = "elasticsearch_connection"
 
     instrument_list = [
-        ("elasticsearch.connection.http_urllib3", "Urllib3HttpConnection.perform_request"),
-        ("elasticsearch.connection.http_requests", "RequestsHttpConnection.perform_request"),
+        ("elasticsearch.transport", "Transport.perform_request"),
     ]
 
     def call(self, module, method, wrapped, instance, args, kwargs):
-        signature = self.get_signature(args, kwargs)
-        context = self.get_context(instance, args, kwargs)
-
         with elasticapm.capture_span(
-            signature,
+            self._get_signature(args, kwargs),
             span_type="db",
             span_subtype="elasticsearch",
             span_action="query",
-            extra=context,
+            extra={},
             skip_frames=2,
             leaf=True,
-        ):
-            return wrapped(*args, **kwargs)
+        ) as span:
+            result_data = wrapped(*args, **kwargs)
+
+            try:
+                span.context["db"]["rows_affected"] = result_data["hits"]["total"]["value"]
+            except (KeyError, TypeError):
+                pass
+
+            return result_data
+
+    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)
diff --git a/elasticapm/instrumentation/register.py b/elasticapm/instrumentation/register.py
--- a/elasticapm/instrumentation/register.py
+++ b/elasticapm/instrumentation/register.py
@@ -54,6 +54,7 @@
     "elasticapm.instrumentation.packages.sqlite.SQLiteInstrumentation",
     "elasticapm.instrumentation.packages.urllib3.Urllib3Instrumentation",
     "elasticapm.instrumentation.packages.elasticsearch.ElasticsearchConnectionInstrumentation",
+    "elasticapm.instrumentation.packages.elasticsearch.ElasticsearchTransportInstrumentation",
     "elasticapm.instrumentation.packages.cassandra.CassandraInstrumentation",
     "elasticapm.instrumentation.packages.pymssql.PyMSSQLInstrumentation",
     "elasticapm.instrumentation.packages.pyodbc.PyODBCInstrumentation",
@@ -73,6 +74,7 @@
             "elasticapm.instrumentation.packages.asyncio.aiohttp_client.AioHttpClientInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.httpx.HttpxAsyncClientInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.elasticsearch.ElasticSearchAsyncConnection",
+            "elasticapm.instrumentation.packages.asyncio.elasticsearch.ElasticsearchAsyncTransportInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.aiopg.AioPGInstrumentation",
             "elasticapm.instrumentation.packages.asyncio.asyncpg.AsyncPGInstrumentation",
             "elasticapm.instrumentation.packages.tornado.TornadoRequestExecuteInstrumentation",


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
@@ -91,6 +91,7 @@ async def test_info(instrument, elasticapm_client, async_elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["sync"] is False
+    assert span["context"]["http"]["status_code"] == 200
 
 
 async def test_create(instrument, elasticapm_client, async_elasticsearch):
@@ -124,6 +125,7 @@ async def test_create(instrument, elasticapm_client, async_elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 201
 
 
 async def test_search_body(instrument, elasticapm_client, async_elasticsearch):
@@ -152,6 +154,9 @@ async def test_search_body(instrument, elasticapm_client, async_elasticsearch):
         '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
     )
     assert span["sync"] is False
+    if ES_VERSION[0] >= 7:
+        assert span["context"]["db"]["rows_affected"] == 1
+    assert span["context"]["http"]["status_code"] == 200
 
 
 async def test_count_body(instrument, elasticapm_client, async_elasticsearch):
@@ -175,3 +180,4 @@ async def test_count_body(instrument, elasticapm_client, async_elasticsearch):
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query": {"term": {"user": "kimchy"}}}')
     assert span["sync"] is False
+    assert span["context"]["http"]["status_code"] == 200
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
@@ -108,6 +108,7 @@ def test_ping(instrument, elasticapm_client, elasticsearch):
         "address": os.environ["ES_URL"],
         "service": {"name": "elasticsearch", "resource": "elasticsearch", "type": "db"},
     }
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -125,6 +126,7 @@ def test_info(instrument, elasticapm_client, elasticsearch):
     assert span["type"] == "db"
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -157,6 +159,7 @@ def test_create(instrument, elasticapm_client, elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 201
 
 
 @pytest.mark.integrationtest
@@ -183,6 +186,7 @@ def test_index(instrument, elasticapm_client, elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 201
 
 
 @pytest.mark.integrationtest
@@ -204,6 +208,7 @@ def test_exists(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.skipif(ES_VERSION[0] < 5, reason="unsupported method")
@@ -232,6 +237,7 @@ def test_exists_source(instrument, elasticapm_client, elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -263,6 +269,7 @@ def test_get(instrument, elasticapm_client, elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -293,6 +300,7 @@ def test_get_source(instrument, elasticapm_client, elasticsearch):
         assert span["action"] == "query"
         assert span["context"]["db"]["type"] == "elasticsearch"
         assert "statement" not in span["context"]["db"]
+        assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -319,6 +327,7 @@ def test_update_document(instrument, elasticapm_client, elasticsearch):
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert "statement" not in span["context"]["db"]
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -347,6 +356,9 @@ def test_search_body(instrument, elasticapm_client, elasticsearch):
     ) or json.loads(span["context"]["db"]["statement"]) == json.loads(
         '{"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}'
     )
+    if ES_VERSION[0] >= 7:
+        assert span["context"]["db"]["rows_affected"] == 1
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -372,6 +384,9 @@ def test_search_querystring(instrument, elasticapm_client, elasticsearch):
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert span["context"]["db"]["statement"] == "q=user:kimchy"
+    if ES_VERSION[0] >= 7:
+        assert span["context"]["db"]["rows_affected"] == 1
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -399,6 +414,7 @@ def test_search_both(instrument, elasticapm_client, elasticsearch):
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert span["context"]["db"]["statement"].startswith('q=text:hola\n\n{"query":')
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -425,6 +441,7 @@ def test_count_body(instrument, elasticapm_client, elasticsearch):
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query": {"term": {"user": "kimchy"}}}')
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -450,6 +467,7 @@ def test_count_querystring(instrument, elasticapm_client, elasticsearch):
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
     assert span["context"]["db"]["statement"] == "q=user:kimchy"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -470,6 +488,7 @@ def test_delete(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -491,6 +510,7 @@ def test_multiple_indexes(instrument, elasticapm_client, elasticsearch):
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.skipif(ES_VERSION[0] >= 7, reason="doc_type unsupported")
@@ -511,6 +531,7 @@ def test_multiple_indexes_doctypes(instrument, elasticapm_client, elasticsearch)
     assert span["subtype"] == "elasticsearch"
     assert span["action"] == "query"
     assert span["context"]["db"]["type"] == "elasticsearch"
+    assert span["context"]["http"]["status_code"] == 200
 
 
 @pytest.mark.integrationtest
@@ -528,3 +549,4 @@ def test_custom_serializer(instrument, elasticapm_client, elasticsearch):
     spans = elasticapm_client.spans_for_transaction(transaction)
     span = spans[0]
     assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query":{"term":{"2":{"value":1}}}}')
+    assert span["context"]["http"]["status_code"] == 200