Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 35 additions & 15 deletions ddtrace/contrib/dbapi/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,15 @@

from ddtrace import Pin
from ddtrace.ext import AppTypes, sql
from ddtrace.settings import config
from ddtrace.utils.formats import asbool, get_env

log = logging.getLogger(__name__)

config._add('dbapi2', dict(
trace_fetch_methods=asbool(get_env('dbapi2', 'trace_fetch_methods', 'false')),
))


class TracedCursor(wrapt.ObjectProxy):
""" TracedCursor wraps a psql cursor and traces it's queries. """
Expand Down Expand Up @@ -72,6 +78,27 @@ def execute(self, query, *args, **kwargs):
self._trace_method(self.__wrapped__.execute, self._self_datadog_name, query, {}, query, *args, **kwargs)
return self

def callproc(self, proc, args):
""" Wraps the cursor.callproc method"""
self._self_last_execute_operation = proc
return self._trace_method(self.__wrapped__.callproc, self._self_datadog_name, proc, {}, proc, args)

def __enter__(self):
# previous versions of the dbapi didn't support context managers. let's
# reference the func that would be called to ensure that errors
# messages will be the same.
self.__wrapped__.__enter__

# and finally, yield the traced cursor.
return self


class FetchTracedCursor(TracedCursor):
"""
Sub-class of :class:`TracedCursor` that also instruments `fetchone`, `fetchall`, and `fetchmany` methods.

We do not trace these functions by default since they can get very noisy (e.g. `fetchone` with 100k rows).
"""
def fetchone(self, *args, **kwargs):
""" Wraps the cursor.fetchone method"""
span_name = '{}.{}'.format(self._self_datadog_name, 'fetchone')
Expand Down Expand Up @@ -101,25 +128,18 @@ def fetchmany(self, *args, **kwargs):
return self._trace_method(self.__wrapped__.fetchmany, span_name, self._self_last_execute_operation, extra_tags,
*args, **kwargs)

def callproc(self, proc, args):
""" Wraps the cursor.callproc method"""
self._self_last_execute_operation = proc
return self._trace_method(self.__wrapped__.callproc, self._self_datadog_name, proc, {}, proc, args)

def __enter__(self):
# previous versions of the dbapi didn't support context managers. let's
# reference the func that would be called to ensure that errors
# messages will be the same.
self.__wrapped__.__enter__

# and finally, yield the traced cursor.
return self


class TracedConnection(wrapt.ObjectProxy):
""" TracedConnection wraps a Connection with tracing code. """

def __init__(self, conn, pin=None, cursor_cls=TracedCursor):
def __init__(self, conn, pin=None, cursor_cls=None):
# Set default cursor class if one was not provided
if not cursor_cls:
# Do not trace `fetch*` methods by default
cursor_cls = TracedCursor
if config.dbapi2.trace_fetch_methods:
cursor_cls = FetchTracedCursor

super(TracedConnection, self).__init__(conn)
name = _get_vendor(conn)
self._self_datadog_name = '{}.connection'.format(name)
Expand Down
19 changes: 13 additions & 6 deletions ddtrace/contrib/psycopg/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import wrapt

# project
from ddtrace import Pin
from ddtrace import Pin, config
from ddtrace.contrib import dbapi
from ddtrace.ext import sql, net, db

Expand Down Expand Up @@ -51,14 +51,21 @@ def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
return super(Psycopg2TracedCursor, self)._trace_method(method, name, resource, extra_tags, *args, **kwargs)


class Psycopg2FetchTracedCursor(Psycopg2TracedCursor, dbapi.FetchTracedCursor):
""" FetchTracedCursor for psycopg2 """


class Psycopg2TracedConnection(dbapi.TracedConnection):
""" TracedConnection wraps a Connection with tracing code. """

def __init__(self, conn, pin=None, cursor_cls=Psycopg2TracedCursor):
super(Psycopg2TracedConnection, self).__init__(conn, pin)
# wrapt requires prefix of `_self` for attributes that are only in the
# proxy (since some of our source objects will use `__slots__`)
self._self_cursor_cls = cursor_cls
def __init__(self, conn, pin=None, cursor_cls=None):
if not cursor_cls:
# Do not trace `fetch*` methods by default
cursor_cls = Psycopg2TracedCursor
if config.dbapi2.trace_fetch_methods:
cursor_cls = Psycopg2FetchTracedCursor

super(Psycopg2TracedConnection, self).__init__(conn, pin, cursor_cls=cursor_cls)


def patch_conn(conn, traced_conn_cls=Psycopg2TracedConnection):
Expand Down
Loading