trace.py

from functools import wraps

import jaeger_client
from flask import current_app
from flask_opentracing import FlaskTracer
from opentracing_instrumentation.client_hooks import requests, strict_redis, sqlalchemy
from opentracing_instrumentation.request_context import RequestContextManager


def install_tracer(service_name, flask_app=None):
    config = jaeger_client.Config(
        config={
            'sampler': {
                'type': 'const',
                'param': 1,
            },
            'logging': True,
        },
        service_name=service_name,
        validate=True)
    tracer = config.initialize_tracer()
    _install_patches()

    if flask_app is not None:
        _install_flask_tracer(flask_app, tracer)


def _install_patches():
    requests.install_patches()
    strict_redis.install_patches()
    sqlalchemy.install_patches()


def _install_flask_tracer(flask_app, tracer):
    flask_app.tracer = FlaskTracer(tracer, True, flask_app)


def controller_action(action):
    @wraps(action)
    def traced_action(*args, **kwds):
        if hasattr(current_app, 'tracer'):
            with RequestContextManager(span=current_app.tracer.get_span()):
                return action(*args, **kwds)
        else:
            return action(*args, **kwds)

    return traced_action