ddtrace flask middleware fix

middleware.py

"""
This code is based on https://github.com/DataDog/dd-trace-py/blob/3f7fbe920d8d8195e4028e522cbe43a974fe91f2/ddtrace/contrib/flask/middleware.py

Copyright (c) 2016, Datadog <info@datadoghq.com>
"""
import logging

from ddtrace import compat
from ddtrace.ext import http, errors, AppTypes
from ddtrace.propagation.http import HTTPPropagator
from flask import g, request, signals

log = logging.getLogger(__name__)


class TraceMiddleware(object):
    def __init__(self, app, tracer, service="flask", distributed_tracing=True):
        self.app = app
        self.app.logger.info("initializing trace middleware")

        # save our traces.
        self._tracer = tracer
        self._service = service
        self._use_distributed_tracing = distributed_tracing

        self._tracer.set_service_info(
            service=service,
            app="flask",
            app_type=AppTypes.web,
        )

        # our signal receivers
        self._receivers = []

        # instrument request timings
        timing_signals = {
            'request_started': self._request_started,
            'got_request_exception': self._got_request_exception,
            'request_finished': self._request_finished,
        }
        self._connect(timing_signals)

    def _connect(self, signal_to_handler):
        connected = True
        for name, handler in signal_to_handler.items():
            s = getattr(signals, name, None)
            if not s:
                connected = False
                log.warning("trying to instrument missing signal %s", name)
                continue
            s.connect(handler, sender=self.app, weak=False)
            self._receivers.append(handler)
        return connected

    def _request_started(self, sender):
        self._start_span()

    def _got_request_exception(self, sender, exception, **kwargs):
        try:
            self._set_exception_info(exception)
        except Exception:
            self.app.logger.exception("error tracing error")

    def _request_finished(self, sender, response, **kwargs):
        try:
            self._finish_span(response)
        except Exception:
            self.app.logger.exception("error finishing trace")
        return response

    def _start_span(self):
        if self._use_distributed_tracing:
            propagator = HTTPPropagator()
            context = propagator.extract(request.headers)
            if context.trace_id:
                self._tracer.context_provider.activate(context)
        try:
            g.flask_datadog_span = self._tracer.trace(
                "flask.request",
                service=self._service,
                span_type=http.TYPE,
            )
        except Exception:
            self.app.logger.exception("error tracing request")

    def _set_exception_info(self, exception):
        span = getattr(g, 'flask_datadog_span', None)
        if span and span.sampled:
            span.set_tag(errors.ERROR_TYPE, type(exception))
            span.set_tag(errors.ERROR_MSG, exception)
            span.set_traceback()

    def _finish_span(self, response):
        span = getattr(g, 'flask_datadog_span', None)
        if span:
            if span.sampled:
                code = response.status_code
                method = request.method if request else None

                span.resource = compat.to_unicode(request.endpoint).lower()
                span.set_tag(http.URL, compat.to_unicode(request.base_url or ''))
                span.set_tag(http.STATUS_CODE, code)
                span.set_tag(http.METHOD, method)
                span.error = int(response.status_code >= 500)
            span.finish()
            g.flask_datadog_span = None