ensure that the elasticsearch instrumentation handles DroppedSpans correctly (elastic#1190)

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"