|
23 | 23 | from opencensus.trace import utils |
24 | 24 | from opencensus.trace.samplers import probability |
25 | 25 |
|
| 26 | +from django.db import connection |
26 | 27 | try: |
27 | 28 | from django.utils.deprecation import MiddlewareMixin |
28 | 29 | except ImportError: # pragma: NO COVER |
@@ -237,9 +238,38 @@ def process_request(self, request): |
237 | 238 | SPAN_THREAD_LOCAL_KEY, |
238 | 239 | span) |
239 | 240 |
|
| 241 | + connection.execute_wrappers.append(self._trace_db_call) |
| 242 | + |
240 | 243 | except Exception: # pragma: NO COVER |
241 | 244 | log.error('Failed to trace request', exc_info=True) |
242 | 245 |
|
| 246 | + def _trace_db_call(self, execute, sql, params, many, context): |
| 247 | + _tracer = execution_context.get_opencensus_tracer() |
| 248 | + if _tracer is not None: |
| 249 | + if many: |
| 250 | + method_name = 'executemany' |
| 251 | + else: |
| 252 | + method_name = 'execute' |
| 253 | + db_type = context['connection'].vendor |
| 254 | + # Note that although get_opencensus_tracer() returns a NoopTracer |
| 255 | + # if no thread local has been set, set_opencensus_tracer() does NOT |
| 256 | + # protect against setting None to the thread local - be defensive |
| 257 | + # here |
| 258 | + _span = _tracer.start_span() |
| 259 | + _span.name = '{}.query'.format(db_type) |
| 260 | + _span.span_kind = span_module.SpanKind.CLIENT |
| 261 | + _tracer.add_attribute_to_current_span( |
| 262 | + '{}.query'.format(db_type), sql) |
| 263 | + _tracer.add_attribute_to_current_span( |
| 264 | + '{}.cursor.method.name'.format(db_type), method_name) |
| 265 | + |
| 266 | + result = execute(sql, params, many, context) |
| 267 | + |
| 268 | + if _tracer is not None: |
| 269 | + _tracer.end_span() |
| 270 | + |
| 271 | + return result |
| 272 | + |
243 | 273 | def process_view(self, request, view_func, *args, **kwargs): |
244 | 274 | """Process view is executed before the view function, here we get the |
245 | 275 | function name add set it as the span name. |
|
0 commit comments