From c440250a7cf1ca5c21d4226b7f938b79fe21b12a Mon Sep 17 00:00:00 2001
From: Leighton Chen <lechen@microsoft.com>
Date: Wed, 14 Oct 2020 15:07:23 -0400
Subject: [PATCH] Use is_recording flag in jinja, celery, esearch, falcon
 instrumentations (#1241)

---
 .../instrumentation/elasticsearch/__init__.py | 45 ++++++++++---------
 .../tests/test_elasticsearch.py               | 18 ++++++++
 2 files changed, 42 insertions(+), 21 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/__init__.py b/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/__init__.py
index f350a7dc2..6e9f411f8 100644
--- a/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/__init__.py
@@ -110,6 +110,7 @@ class ElasticsearchInstrumentor(BaseInstrumentor):
 
 
 def _wrap_perform_request(tracer, span_name_prefix):
+    # pylint: disable=R0912
     def wrapper(wrapped, _, args, kwargs):
         method = url = None
         try:
@@ -125,26 +126,27 @@ def _wrap_perform_request(tracer, span_name_prefix):
         params = kwargs.get("params", {})
         body = kwargs.get("body", None)
 
-        attributes = {
-            "component": "elasticsearch-py",
-            "db.type": "elasticsearch",
-        }
-
-        if url:
-            attributes["elasticsearch.url"] = url
-        if method:
-            attributes["elasticsearch.method"] = method
-        if body:
-            attributes["db.statement"] = str(body)
-        if params:
-            attributes["elasticsearch.params"] = str(params)
-
         with tracer.start_as_current_span(
-            op_name, kind=SpanKind.CLIENT, attributes=attributes
+            op_name, kind=SpanKind.CLIENT,
         ) as span:
+            if span.is_recording():
+                attributes = {
+                    "component": "elasticsearch-py",
+                    "db.type": "elasticsearch",
+                }
+                if url:
+                    attributes["elasticsearch.url"] = url
+                if method:
+                    attributes["elasticsearch.method"] = method
+                if body:
+                    attributes["db.statement"] = str(body)
+                if params:
+                    attributes["elasticsearch.params"] = str(params)
+                for key, value in attributes.items():
+                    span.set_attribute(key, value)
             try:
                 rv = wrapped(*args, **kwargs)
-                if isinstance(rv, dict):
+                if isinstance(rv, dict) and span.is_recording():
                     for member in _ATTRIBUTES_FROM_RESULT:
                         if member in rv:
                             span.set_attribute(
@@ -153,11 +155,12 @@ def _wrap_perform_request(tracer, span_name_prefix):
                             )
                 return rv
             except Exception as ex:  # pylint: disable=broad-except
-                if isinstance(ex, elasticsearch.exceptions.NotFoundError):
-                    status = StatusCanonicalCode.NOT_FOUND
-                else:
-                    status = StatusCanonicalCode.UNKNOWN
-                span.set_status(Status(status, str(ex)))
+                if span.is_recording():
+                    if isinstance(ex, elasticsearch.exceptions.NotFoundError):
+                        status = StatusCanonicalCode.NOT_FOUND
+                    else:
+                        status = StatusCanonicalCode.UNKNOWN
+                    span.set_status(Status(status, str(ex)))
                 raise ex
 
     return wrapper
diff --git a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py
index cc1d31477..3d93838fe 100644
--- a/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py
+++ b/instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py
@@ -88,6 +88,24 @@ class TestElasticsearchIntegration(TestBase):
         spans_list = self.get_ordered_finished_spans()
         self.assertEqual(len(spans_list), 1)
 
+    def test_span_not_recording(self, request_mock):
+        request_mock.return_value = (1, {}, {})
+        mock_tracer = mock.Mock()
+        mock_span = mock.Mock()
+        mock_span.is_recording.return_value = False
+        mock_tracer.start_span.return_value = mock_span
+        mock_tracer.use_span.return_value.__enter__ = mock_span
+        mock_tracer.use_span.return_value.__exit__ = mock_span
+        with mock.patch("opentelemetry.trace.get_tracer") as tracer:
+            tracer.return_value = mock_tracer
+            Elasticsearch()
+            self.assertFalse(mock_span.is_recording())
+            self.assertTrue(mock_span.is_recording.called)
+            self.assertFalse(mock_span.set_attribute.called)
+            self.assertFalse(mock_span.set_status.called)
+
+        ElasticsearchInstrumentor().uninstrument()
+
     def test_prefix_arg(self, request_mock):
         prefix = "prefix-from-env"
         ElasticsearchInstrumentor().uninstrument()