ensure that the elasticsearch instrumentation handles DroppedSpans correctly

If the elasticsearch span is dropped for some reason, the context object
is None, which, if unhandled, leads to an exception

fixes elastic#1188

Author: beniwohli

elasticapm/instrumentation/packages/elasticsearch.py

@@ -34,7 +34,7 @@
 
 import elasticapm
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
-from elasticapm.traces import execution_context
+from elasticapm.traces import DroppedSpan, execution_context
 from elasticapm.utils.logging import get_logger
 
 logger = get_logger("elasticapm.instrument")
@@ -52,6 +52,8 @@ class ElasticsearchConnectionInstrumentation(AbstractInstrumentedModule):
 
     def call(self, module, method, wrapped, instance, args, kwargs):
         span = execution_context.get_span()
+        if isinstance(span, DroppedSpan):
+            return wrapped(*args, **kwargs)
 
         self._update_context_by_request_data(span.context, instance, args, kwargs)
 

tests/instrumentation/elasticsearch_tests.py

@@ -39,6 +39,7 @@
 from elasticsearch import Elasticsearch
 from elasticsearch.serializer import JSONSerializer
 
+import elasticapm
 from elasticapm.conf.constants import TRANSACTION
 from elasticapm.utils import compat
 
@@ -550,3 +551,17 @@ def test_custom_serializer(instrument, elasticapm_client, elasticsearch):
     span = spans[0]
     assert json.loads(span["context"]["db"]["statement"]) == json.loads('{"query":{"term":{"2":{"value":1}}}}')
     assert span["context"]["http"]["status_code"] == 200
+
+
+@pytest.mark.integrationtest
+def test_dropped_span(instrument, elasticapm_client, elasticsearch):
+    elasticapm_client.begin_transaction("test")
+    with elasticapm.capture_span("test", leaf=True):
+        elasticsearch.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"] == "test"