Merge pull request #833 from DataDog/0.21.1-dev

Release 0.21.1

Author: brettlangdon

ddtrace/__init__.py

@@ -4,7 +4,7 @@
 from .tracer import Tracer
 from .settings import config
 
-__version__ = '0.21.0'
+__version__ = '0.21.1'
 
 # a global tracer instance with integration settings
 tracer = Tracer()

ddtrace/bootstrap/sitecustomize.py

@@ -15,6 +15,10 @@
     '[dd.trace_id=%(dd.trace_id)s dd.span_id=%(dd.span_id)s] ' if logs_injection else ''
 )
 
+if logs_injection:
+    # immediately patch logging if trace id injected
+    from ddtrace import patch; patch(logging=True)  # noqa
+
 debug = os.environ.get("DATADOG_TRACE_DEBUG")
 
 # Set here a default logging format for basicConfig

ddtrace/contrib/dbapi/__init__.py

@@ -65,18 +65,23 @@ def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
     def executemany(self, query, *args, **kwargs):
         """ Wraps the cursor.executemany method"""
         self._self_last_execute_operation = query
+        # Always return the result as-is
+        # DEV: Some libraries return `None`, others `int`, and others the cursor objects
+        #      These differences should be overriden at the integration specific layer (e.g. in `sqlite3/patch.py`)
         # FIXME[matt] properly handle kwargs here. arg names can be different
         # with different libs.
-        self._trace_method(
+        return self._trace_method(
             self.__wrapped__.executemany, self._self_datadog_name, query, {'sql.executemany': 'true'},
             query, *args, **kwargs)
-        return self
 
     def execute(self, query, *args, **kwargs):
         """ Wraps the cursor.execute method"""
         self._self_last_execute_operation = query
-        self._trace_method(self.__wrapped__.execute, self._self_datadog_name, query, {}, query, *args, **kwargs)
-        return self
+
+        # Always return the result as-is
+        # DEV: Some libraries return `None`, others `int`, and others the cursor objects
+        #      These differences should be overriden at the integration specific layer (e.g. in `sqlite3/patch.py`)
+        return self._trace_method(self.__wrapped__.execute, self._self_datadog_name, query, {}, query, *args, **kwargs)
 
     def callproc(self, proc, args):
         """ Wraps the cursor.callproc method"""

ddtrace/contrib/sqlite3/patch.py

@@ -4,10 +4,10 @@
 import wrapt
 
 # project
-from ddtrace import Pin
-from ddtrace.contrib.dbapi import TracedConnection
-
+from ...contrib.dbapi import TracedConnection, TracedCursor, FetchTracedCursor
 from ...ext import AppTypes
+from ...pin import Pin
+from ...settings import config
 
 # Original connect method
 _connect = sqlite3.connect
@@ -36,7 +36,31 @@ def patch_conn(conn):
     return wrapped
 
 
+class TracedSQLiteCursor(TracedCursor):
+    def executemany(self, *args, **kwargs):
+        # DEV: SQLite3 Cursor.execute always returns back the cursor instance
+        super(TracedSQLiteCursor, self).executemany(*args, **kwargs)
+        return self
+
+    def execute(self, *args, **kwargs):
+        # DEV: SQLite3 Cursor.execute always returns back the cursor instance
+        super(TracedSQLiteCursor, self).execute(*args, **kwargs)
+        return self
+
+
+class TracedSQLiteFetchCursor(TracedSQLiteCursor, FetchTracedCursor):
+    pass
+
+
 class TracedSQLite(TracedConnection):
+    def __init__(self, conn, pin=None, cursor_cls=None):
+        if not cursor_cls:
+            # Do not trace `fetch*` methods by default
+            cursor_cls = TracedSQLiteCursor
+            if config.dbapi2.trace_fetch_methods:
+                cursor_cls = TracedSQLiteFetchCursor
+
+            super(TracedSQLite, self).__init__(conn, pin=pin, cursor_cls=cursor_cls)
 
     def execute(self, *args, **kwargs):
         # sqlite has a few extra sugar functions

ddtrace/propagation/http.py

@@ -56,34 +56,35 @@ def parent_call():
             headers[HTTP_HEADER_SAMPLING_PRIORITY] = str(span_context.sampling_priority)
 
     @staticmethod
-    def extract_trace_id(headers):
-        trace_id = 0
+    def extract_header_value(possible_header_names, headers, default=None):
+        for header, value in headers.items():
+            for header_name in possible_header_names:
+                if header.lower() == header_name.lower():
+                    return value
 
-        for key in POSSIBLE_HTTP_HEADER_TRACE_IDS:
-            if key in headers:
-                trace_id = headers.get(key)
+        return default
 
-        return int(trace_id)
+    @staticmethod
+    def extract_trace_id(headers):
+        return int(
+            HTTPPropagator.extract_header_value(
+                POSSIBLE_HTTP_HEADER_TRACE_IDS, headers, default=0,
+            )
+        )
 
     @staticmethod
     def extract_parent_span_id(headers):
-        parent_span_id = 0
-
-        for key in POSSIBLE_HTTP_HEADER_PARENT_IDS:
-            if key in headers:
-                parent_span_id = headers.get(key)
-
-        return int(parent_span_id)
+        return int(
+            HTTPPropagator.extract_header_value(
+                POSSIBLE_HTTP_HEADER_PARENT_IDS, headers, default=0,
+            )
+        )
 
     @staticmethod
     def extract_sampling_priority(headers):
-        sampling_priority = None
-
-        for key in POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES:
-            if key in headers:
-                sampling_priority = headers.get(key)
-
-        return sampling_priority
+        return HTTPPropagator.extract_header_value(
+            POSSIBLE_HTTP_HEADER_SAMPLING_PRIORITIES, headers,
+        )
 
     def extract(self, headers):
         """Extract a Context from HTTP headers into a new Context.

tests/contrib/dbapi/test_unit.py

@@ -14,17 +14,23 @@ def setUp(self):
     def test_execute_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0
+        cursor.execute.return_value = '__result__'
+
         pin = Pin('pin_name', tracer=self.tracer)
         traced_cursor = TracedCursor(cursor, pin)
-        assert traced_cursor is traced_cursor.execute('__query__', 'arg_1', kwarg1='kwarg1')
+        # DEV: We always pass through the result
+        assert '__result__' == traced_cursor.execute('__query__', 'arg_1', kwarg1='kwarg1')
         cursor.execute.assert_called_once_with('__query__', 'arg_1', kwarg1='kwarg1')
 
     def test_executemany_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0
+        cursor.executemany.return_value = '__result__'
+
         pin = Pin('pin_name', tracer=self.tracer)
         traced_cursor = TracedCursor(cursor, pin)
-        assert traced_cursor is traced_cursor.executemany('__query__', 'arg_1', kwarg1='kwarg1')
+        # DEV: We always pass through the result
+        assert '__result__' == traced_cursor.executemany('__query__', 'arg_1', kwarg1='kwarg1')
         cursor.executemany.assert_called_once_with('__query__', 'arg_1', kwarg1='kwarg1')
 
     def test_fetchone_wrapped_is_called_and_returned(self):
@@ -114,14 +120,17 @@ def test_when_pin_disabled_then_no_tracing(self):
         cursor = self.cursor
         tracer = self.tracer
         cursor.rowcount = 0
+        cursor.execute.return_value = '__result__'
+        cursor.executemany.return_value = '__result__'
+
         tracer.enabled = False
         pin = Pin('pin_name', tracer=tracer)
         traced_cursor = TracedCursor(cursor, pin)
 
-        assert traced_cursor is traced_cursor.execute('arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.execute('arg_1', kwarg1='kwarg1')
         assert len(tracer.writer.pop()) == 0
 
-        assert traced_cursor is traced_cursor.executemany('arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.executemany('arg_1', kwarg1='kwarg1')
         assert len(tracer.writer.pop()) == 0
 
         cursor.callproc.return_value = 'callproc'
@@ -191,17 +200,21 @@ def setUp(self):
     def test_execute_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0
+        cursor.execute.return_value = '__result__'
+
         pin = Pin('pin_name', tracer=self.tracer)
         traced_cursor = FetchTracedCursor(cursor, pin)
-        assert traced_cursor is traced_cursor.execute('__query__', 'arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.execute('__query__', 'arg_1', kwarg1='kwarg1')
         cursor.execute.assert_called_once_with('__query__', 'arg_1', kwarg1='kwarg1')
 
     def test_executemany_wrapped_is_called_and_returned(self):
         cursor = self.cursor
         cursor.rowcount = 0
+        cursor.executemany.return_value = '__result__'
+
         pin = Pin('pin_name', tracer=self.tracer)
         traced_cursor = FetchTracedCursor(cursor, pin)
-        assert traced_cursor is traced_cursor.executemany('__query__', 'arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.executemany('__query__', 'arg_1', kwarg1='kwarg1')
         cursor.executemany.assert_called_once_with('__query__', 'arg_1', kwarg1='kwarg1')
 
     def test_fetchone_wrapped_is_called_and_returned(self):
@@ -297,14 +310,17 @@ def test_when_pin_disabled_then_no_tracing(self):
         cursor = self.cursor
         tracer = self.tracer
         cursor.rowcount = 0
+        cursor.execute.return_value = '__result__'
+        cursor.executemany.return_value = '__result__'
+
         tracer.enabled = False
         pin = Pin('pin_name', tracer=tracer)
         traced_cursor = FetchTracedCursor(cursor, pin)
 
-        assert traced_cursor is traced_cursor.execute('arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.execute('arg_1', kwarg1='kwarg1')
         assert len(tracer.writer.pop()) == 0
 
-        assert traced_cursor is traced_cursor.executemany('arg_1', kwarg1='kwarg1')
+        assert '__result__' == traced_cursor.executemany('arg_1', kwarg1='kwarg1')
         assert len(tracer.writer.pop()) == 0
 
         cursor.callproc.return_value = 'callproc'

tests/contrib/mysqldb/test_mysql.py

@@ -42,7 +42,8 @@ def test_simple_query(self):
         conn, tracer = self._get_conn_tracer()
         writer = tracer.writer
         cursor = conn.cursor()
-        cursor.execute("SELECT 1")
+        rowcount = cursor.execute("SELECT 1")
+        eq_(rowcount, 1)
         rows = cursor.fetchall()
         eq_(len(rows), 1)
         spans = writer.pop()

tests/contrib/psycopg/test_psycopg.py

@@ -89,7 +89,8 @@ def assert_conn_is_traced(self, db, service):
 
         start = time.time()
         cursor = db.cursor()
-        cursor.execute(q)
+        res = cursor.execute(q)
+        self.assertIsNone(res)
         rows = cursor.fetchall()
         end = time.time()
 

tests/contrib/pymysql/test_pymysql.py

@@ -54,7 +54,11 @@ def test_simple_query(self):
         conn, tracer = self._get_conn_tracer()
         writer = tracer.writer
         cursor = conn.cursor()
-        cursor.execute('SELECT 1')
+
+        # PyMySQL returns back the rowcount instead of a cursor
+        rowcount = cursor.execute('SELECT 1')
+        eq_(rowcount, 1)
+
         rows = cursor.fetchall()
         eq_(len(rows), 1)
         spans = writer.pop()
@@ -135,7 +139,11 @@ def test_query_many(self):
         stmt = "INSERT INTO dummy (dummy_key, dummy_value) VALUES (%s, %s)"
         data = [("foo", "this is foo"),
                 ("bar", "this is bar")]
-        cursor.executemany(stmt, data)
+
+        # PyMySQL `executemany()` returns the rowcount
+        rowcount = cursor.executemany(stmt, data)
+        eq_(rowcount, 2)
+
         query = "SELECT dummy_key, dummy_value FROM dummy ORDER BY dummy_key"
         cursor.execute(query)
         rows = cursor.fetchall()

tests/contrib/sqlite3/test_sqlite3.py

@@ -6,7 +6,7 @@
 import ddtrace
 from ddtrace import Pin
 from ddtrace.contrib.sqlite3 import connection_factory
-from ddtrace.contrib.sqlite3.patch import patch, unpatch
+from ddtrace.contrib.sqlite3.patch import patch, unpatch, TracedSQLiteCursor
 from ddtrace.ext import errors
 
 # testing
@@ -29,8 +29,9 @@ def test_backwards_compat(self):
         factory = connection_factory(self.tracer, service='my_db_service')
         conn = sqlite3.connect(':memory:', factory=factory)
         q = 'select * from sqlite_master'
-        rows = conn.execute(q)
-        assert not rows.fetchall()
+        cursor = conn.execute(q)
+        self.assertIsInstance(cursor, TracedSQLiteCursor)
+        assert not cursor.fetchall()
         assert not self.spans
 
     def test_service_info(self):
@@ -60,6 +61,7 @@ def test_sqlite(self):
             q = 'select * from sqlite_master'
             start = time.time()
             cursor = db.execute(q)
+            self.assertIsInstance(cursor, TracedSQLiteCursor)
             rows = cursor.fetchall()
             end = time.time()
             assert not rows