Disable analytics for fetch* and connections

Author: majorgreys

ddtrace/contrib/dbapi/__init__.py

@@ -50,11 +50,12 @@ def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
             s.set_tags(pin.tags)
             s.set_tags(extra_tags)
 
-            # set analytics sample rate if enabled
-            s.set_tag(
-                ANALYTICS_SAMPLE_RATE_KEY,
-                config.dbapi2.get_analytics_sample_rate()
-            )
+            # set analytics sample rate if enabled but only for non-FetchTracedCursor
+            if type(self) is not FetchTracedCursor:
+                s.set_tag(
+                    ANALYTICS_SAMPLE_RATE_KEY,
+                    config.dbapi2.get_analytics_sample_rate()
+                )
 
             try:
                 return method(*args, **kwargs)
@@ -170,12 +171,6 @@ def _trace_method(self, method, name, extra_tags, *args, **kwargs):
             s.set_tags(pin.tags)
             s.set_tags(extra_tags)
 
-            # set analytics sample rate if enabled
-            s.set_tag(
-                ANALYTICS_SAMPLE_RATE_KEY,
-                config.dbapi2.get_analytics_sample_rate()
-            )
-
             return method(*args, **kwargs)
 
     def cursor(self, *args, **kwargs):

tests/contrib/dbapi/test_unit.py

@@ -415,6 +415,41 @@ def method():
         assert span.get_metric('db.rowcount') == 123, 'Row count is set as a metric'
         assert span.get_tag('sql.rows') == '123', 'Row count is set as a tag (for legacy django cursor replacement)'
 
+    def test_fetch_analytics_off(self):
+        with self.override_config(
+                'dbapi2',
+                dict(analytics_enabled=True)
+        ):
+            cursor = self.cursor
+            cursor.rowcount = 0
+            cursor.fetchone.return_value = '__result__'
+            pin = Pin('pin_name', tracer=self.tracer)
+            traced_cursor = FetchTracedCursor(cursor, pin)
+            assert '__result__' == traced_cursor.fetchone('arg_1', kwarg1='kwarg1')
+
+            span = self.tracer.writer.pop()[0]
+            self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))
+
+            cursor = self.cursor
+            cursor.rowcount = 0
+            cursor.fetchall.return_value = '__result__'
+            pin = Pin('pin_name', tracer=self.tracer)
+            traced_cursor = FetchTracedCursor(cursor, pin)
+            assert '__result__' == traced_cursor.fetchall('arg_1', kwarg1='kwarg1')
+
+            span = self.tracer.writer.pop()[0]
+            self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))
+
+            cursor = self.cursor
+            cursor.rowcount = 0
+            cursor.fetchmany.return_value = '__result__'
+            pin = Pin('pin_name', tracer=self.tracer)
+            traced_cursor = FetchTracedCursor(cursor, pin)
+            assert '__result__' == traced_cursor.fetchmany('arg_1', kwarg1='kwarg1')
+
+            span = self.tracer.writer.pop()[0]
+            self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))
+
 
 class TestTracedConnection(BaseTracerTestCase):
     def setUp(self):
@@ -458,7 +493,7 @@ def test_rollback_is_traced(self):
         assert tracer.writer.pop()[0].name == 'mock.connection.rollback'
         connection.rollback.assert_called_with()
 
-    def test_cursor_analytics_with_rate(self):
+    def test_connection_analytics_with_rate(self):
         with self.override_config(
                 'dbapi2',
                 dict(analytics_enabled=True, analytics_sample_rate=0.5)
@@ -470,4 +505,4 @@ def test_cursor_analytics_with_rate(self):
             traced_connection = TracedConnection(connection, pin)
             traced_connection.commit()
             span = tracer.writer.pop()[0]
-            self.assertEqual(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY), 0.5)
+            self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))