|
8 | 8 |
|
9 | 9 | from ddtrace import Pin |
10 | 10 | from ddtrace.ext import AppTypes, sql |
| 11 | +from ddtrace.settings import config |
| 12 | +from ddtrace.utils.formats import asbool, get_env |
11 | 13 |
|
12 | 14 | log = logging.getLogger(__name__) |
13 | 15 |
|
| 16 | +config._add('dbapi2', dict( |
| 17 | + trace_fetch_methods=asbool(get_env('dbapi2', 'trace_fetch_methods', 'false')), |
| 18 | +)) |
| 19 | + |
14 | 20 |
|
15 | 21 | class TracedCursor(wrapt.ObjectProxy): |
16 | 22 | """ TracedCursor wraps a psql cursor and traces it's queries. """ |
@@ -72,6 +78,27 @@ def execute(self, query, *args, **kwargs): |
72 | 78 | self._trace_method(self.__wrapped__.execute, self._self_datadog_name, query, {}, query, *args, **kwargs) |
73 | 79 | return self |
74 | 80 |
|
| 81 | + def callproc(self, proc, args): |
| 82 | + """ Wraps the cursor.callproc method""" |
| 83 | + self._self_last_execute_operation = proc |
| 84 | + return self._trace_method(self.__wrapped__.callproc, self._self_datadog_name, proc, {}, proc, args) |
| 85 | + |
| 86 | + def __enter__(self): |
| 87 | + # previous versions of the dbapi didn't support context managers. let's |
| 88 | + # reference the func that would be called to ensure that errors |
| 89 | + # messages will be the same. |
| 90 | + self.__wrapped__.__enter__ |
| 91 | + |
| 92 | + # and finally, yield the traced cursor. |
| 93 | + return self |
| 94 | + |
| 95 | + |
| 96 | +class FetchTracedCursor(TracedCursor): |
| 97 | + """ |
| 98 | + Sub-class of :class:`TracedCursor` that also instruments `fetchone`, `fetchall`, and `fetchmany` methods. |
| 99 | +
|
| 100 | + We do not trace these functions by default since they can get very noisy (e.g. `fetchone` with 100k rows). |
| 101 | + """ |
75 | 102 | def fetchone(self, *args, **kwargs): |
76 | 103 | """ Wraps the cursor.fetchone method""" |
77 | 104 | span_name = '{}.{}'.format(self._self_datadog_name, 'fetchone') |
@@ -101,25 +128,18 @@ def fetchmany(self, *args, **kwargs): |
101 | 128 | return self._trace_method(self.__wrapped__.fetchmany, span_name, self._self_last_execute_operation, extra_tags, |
102 | 129 | *args, **kwargs) |
103 | 130 |
|
104 | | - def callproc(self, proc, args): |
105 | | - """ Wraps the cursor.callproc method""" |
106 | | - self._self_last_execute_operation = proc |
107 | | - return self._trace_method(self.__wrapped__.callproc, self._self_datadog_name, proc, {}, proc, args) |
108 | | - |
109 | | - def __enter__(self): |
110 | | - # previous versions of the dbapi didn't support context managers. let's |
111 | | - # reference the func that would be called to ensure that errors |
112 | | - # messages will be the same. |
113 | | - self.__wrapped__.__enter__ |
114 | | - |
115 | | - # and finally, yield the traced cursor. |
116 | | - return self |
117 | | - |
118 | 131 |
|
119 | 132 | class TracedConnection(wrapt.ObjectProxy): |
120 | 133 | """ TracedConnection wraps a Connection with tracing code. """ |
121 | 134 |
|
122 | | - def __init__(self, conn, pin=None, cursor_cls=TracedCursor): |
| 135 | + def __init__(self, conn, pin=None, cursor_cls=None): |
| 136 | + # Set default cursor class if one was not provided |
| 137 | + if not cursor_cls: |
| 138 | + # Do not trace `fetch*` methods by default |
| 139 | + cursor_cls = TracedCursor |
| 140 | + if config.dbapi2.trace_fetch_methods: |
| 141 | + cursor_cls = FetchTracedCursor |
| 142 | + |
123 | 143 | super(TracedConnection, self).__init__(conn) |
124 | 144 | name = _get_vendor(conn) |
125 | 145 | self._self_datadog_name = '{}.connection'.format(name) |
|
0 commit comments