From dbc5f3b1667daf487e04293a15fb34ce39ac6bc4 Mon Sep 17 00:00:00 2001
From: Leighton Chen <lechen@microsoft.com>
Date: Thu, 8 Oct 2020 15:25:20 -0400
Subject: [PATCH] Use is_recording flag in aiopg, asyncpg, dbapi, psycopg2,
 pymemcache, pymongo, redis, sqlalchemy instrumentations (#1212)

---
 .../instrumentation/sqlalchemy/engine.py      | 49 +++++++++++--------
 .../tests/test_sqlalchemy.py                  | 23 +++++++++
 2 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
index df80c4841..83a5b82b2 100644
--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py
@@ -79,15 +79,16 @@ class EngineTracer:
     def _before_cur_exec(self, conn, cursor, statement, *args):
         self.current_span = self.tracer.start_span(self.name)
         with self.tracer.use_span(self.current_span, end_on_exit=False):
-            self.current_span.set_attribute("service", self.vendor)
-            self.current_span.set_attribute(_STMT, statement)
+            if self.current_span.is_recording():
+                self.current_span.set_attribute("service", self.vendor)
+                self.current_span.set_attribute(_STMT, statement)
 
-            if not _set_attributes_from_url(
-                self.current_span, conn.engine.url
-            ):
-                _set_attributes_from_cursor(
-                    self.current_span, self.vendor, cursor
-                )
+                if not _set_attributes_from_url(
+                    self.current_span, conn.engine.url
+                ):
+                    _set_attributes_from_cursor(
+                        self.current_span, self.vendor, cursor
+                    )
 
     # pylint: disable=unused-argument
     def _after_cur_exec(self, conn, cursor, statement, *args):
@@ -95,7 +96,11 @@ class EngineTracer:
             return
 
         try:
-            if cursor and cursor.rowcount >= 0:
+            if (
+                cursor
+                and cursor.rowcount >= 0
+                and self.current_span.is_recording()
+            ):
                 self.current_span.set_attribute(_ROWS, cursor.rowcount)
         finally:
             self.current_span.end()
@@ -105,30 +110,34 @@ class EngineTracer:
             return
 
         try:
-            self.current_span.set_status(
-                Status(
-                    StatusCanonicalCode.UNKNOWN,
-                    str(context.original_exception),
+            if self.current_span.is_recording():
+                self.current_span.set_status(
+                    Status(
+                        StatusCanonicalCode.UNKNOWN,
+                        str(context.original_exception),
+                    )
                 )
-            )
         finally:
             self.current_span.end()
 
 
 def _set_attributes_from_url(span: trace.Span, url):
     """Set connection tags from the url. return true if successful."""
-    if url.host:
-        span.set_attribute(_HOST, url.host)
-    if url.port:
-        span.set_attribute(_PORT, url.port)
-    if url.database:
-        span.set_attribute(_DB, url.database)
+    if span.is_recording():
+        if url.host:
+            span.set_attribute(_HOST, url.host)
+        if url.port:
+            span.set_attribute(_PORT, url.port)
+        if url.database:
+            span.set_attribute(_DB, url.database)
 
     return bool(url.host)
 
 
 def _set_attributes_from_cursor(span: trace.Span, vendor, cursor):
     """Attempt to set db connection attributes by introspecting the cursor."""
+    if not span.is_recording():
+        return
     if vendor == "postgres":
         # pylint: disable=import-outside-toplevel
         from psycopg2.extensions import parse_dsn
diff --git a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
index 06593da94..3b2e3548a 100644
--- a/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
+++ b/instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from unittest import mock
 
 from sqlalchemy import create_engine
 
@@ -37,6 +38,28 @@ class TestSqlalchemyInstrumentation(TestBase):
         self.assertEqual(len(spans), 1)
         self.assertEqual(spans[0].name, "sqlite.query")
 
+    def test_not_recording(self):
+        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__ = True
+        with mock.patch("opentelemetry.trace.get_tracer") as tracer:
+            tracer.return_value = mock_tracer
+            engine = create_engine("sqlite:///:memory:")
+            SQLAlchemyInstrumentor().instrument(
+                engine=engine,
+                tracer_provider=self.tracer_provider,
+                service="my-database",
+            )
+            cnx = engine.connect()
+            cnx.execute("SELECT	1 + 1;").fetchall()
+            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)
+
     def test_create_engine_wrapper(self):
         SQLAlchemyInstrumentor().instrument()
         from sqlalchemy import create_engine  # pylint: disable-all