Handle track_total_hits=False in elasticsearch instrumentation (#1687)

* Handle track_total_hits=False

* CHANGELOG

* Fix test

* Fix test (again)

Author: basepi

CHANGELOG.asciidoc

@@ -29,6 +29,18 @@ endif::[]
 //===== Bug fixes
 //
 
+=== Unreleased
+
+// Unreleased changes go here
+// When the next release happens, nest these changes under the "Python Agent version 6.x" heading
+//[float]
+//===== Features
+//
+[float]
+===== Bug fixes
+
+* Fix elasticsearch for track_total_hits=False {pull}1687[#1687]
+
 [[release-notes-6.x]]
 === Python Agent version 6.x
 

elasticapm/instrumentation/packages/elasticsearch.py

@@ -177,8 +177,8 @@ def _get_signature(self, args, kwargs):
 
     def _get_hits(self, result) -> Optional[int]:
         if getattr(result, "body", None) and "hits" in result.body:  # ES >= 8
-            return result.body["hits"]["total"]["value"]
-        elif isinstance(result, dict) and "hits" in result:
+            return result.body["hits"].get("total", {}).get("value")
+        elif isinstance(result, dict) and "hits" in result and "total" in result["hits"]:
             return (
                 result["hits"]["total"]["value"]
                 if isinstance(result["hits"]["total"], dict)

tests/instrumentation/elasticsearch_tests.py

@@ -399,6 +399,32 @@ def test_search_body(instrument, elasticapm_client, elasticsearch):
     assert span["context"]["http"]["status_code"] == 200
 
 
+@pytest.mark.skipif(ES_VERSION[0] < 7, reason="track_total_hits unsupported")
+@pytest.mark.integrationtest
+def test_search_track_total_hits_false(instrument, elasticapm_client, elasticsearch):
+    elasticsearch.create(
+        index="tweets", id="1", refresh=True, **get_kwargs({"user": "kimchy", "text": "hola", "userid": 1})
+    )
+    elasticapm_client.begin_transaction("test")
+    search_query = {"query": {"term": {"user": "kimchy"}}, "sort": ["userid"]}
+    result = elasticsearch.search(body=search_query, track_total_hits=False)
+    elasticapm_client.end_transaction("test", "OK")
+
+    transaction = elasticapm_client.events[TRANSACTION][0]
+    assert result["hits"]["hits"][0]["_source"] == {"user": "kimchy", "text": "hola", "userid": 1}
+    spans = elasticapm_client.spans_for_transaction(transaction)
+    assert len(spans) == 1
+    span = spans[0]
+    # Depending on ES_VERSION, could be /_all/_search or /_search, and GET or POST
+    assert span["name"] in ("ES GET /_search", "ES GET /_all/_search", "ES POST /_search")
+    assert span["type"] == "db"
+    assert span["subtype"] == "elasticsearch"
+    assert span["action"] == "query"
+    assert span["context"]["db"]["type"] == "elasticsearch"
+    assert "rows_affected" not in span["context"]["db"]
+    assert span["context"]["http"]["status_code"] == 200
+
+
 @pytest.mark.integrationtest
 def test_search_querystring(instrument, elasticapm_client, elasticsearch):
     elasticsearch.create(index="tweets", id="1", refresh=True, **get_kwargs({"user": "kimchy", "text": "hola"}))