django_correlation_sql.py

from django.db import connection

from opencensus.trace.execution_context import get_opencensus_tracer


class SQLCorrelationWrapper:
    def __init__(self, trace_id):
        self.trace_id = trace_id

    def __call__(self, execute, sql, params, many, context):
        if len(sql) > 6:
            operation = sql[:6]
            if operation in ('SELECT', 'INSERT', 'DELETE', 'UPDATE'):
                sql = operation + r' /* {} */'.format(self.trace_id) + sql[6:]
                
        execute(sql, params, many, context)


class SQLCorrelationMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        trace_id = get_opencensus_tracer().span_context.trace_id
        with connection.execute_wrapper(SQLCorrelationWrapper(trace_id)):
            response = self.get_response(request)

        return response