Fix lost db context in elasticsearch instrumentation + Elastic Cloud (#1878)

* Fix instance information overriding db context in elasticsearch

* CHANGELOG

Author: basepi

CHANGELOG.asciidoc

@@ -38,9 +38,10 @@ endif::[]
 
 * Add support for grpc aio server interceptor {pull}1870[#1870]
 
-//[float]
-//===== Bug fixes
-//
+[float]
+===== Bug fixes
+
+* Fix a bug in the Elasticsearch client instrumentation which was causing loss of database context (including statement) when interacting with Elastic Cloud {pull}1878[#1878]
 
 
 [[release-notes-6.x]]

elasticapm/instrumentation/packages/asyncio/elasticsearch.py

@@ -63,10 +63,14 @@ async def call(self, module, method, wrapped, instance, args, kwargs):
         result = await wrapped(*args, **kwargs)
         if hasattr(result, "meta"):  # elasticsearch-py 8.x+
             status_code = result.meta.status
+            cluster = result.meta.headers.get("x-found-handling-cluster")
         else:
             status_code = result[0]
+            cluster = result[1].get("x-found-handling-cluster")
 
-        span.context["http"] = {"status_code": status_code}
+        span.update_context("http", {"status_code": status_code})
+        if cluster:
+            span.update_context("db", {"instance": cluster})
 
         return result
 

elasticapm/instrumentation/packages/elasticsearch.py

@@ -75,9 +75,9 @@ def call(self, module, method, wrapped, instance, args, kwargs):
         else:
             status_code = result[0]
             cluster = result[1].get("x-found-handling-cluster")
-        span.context["http"] = {"status_code": status_code}
+        span.update_context("http", {"status_code": status_code})
         if cluster:
-            span.context["db"] = {"instance": cluster}
+            span.update_context("db", {"instance": cluster})
 
         return result