stats middleware

gistfile1.txt

# -*- coding: utf-8 -*-

import time


class StatsdMiddleware(object):
    """
    WSGI middleware that tracks how long a request takes and emits to a statsd client
    """
    def __init__(self,
                 wsgi_app,
                 flask_app,
                 statsd,
                 prefix=None,
                 blacklist=[]):
        self.flask_app = flask_app
        self.wsgi_app = wsgi_app
        self.statsd = statsd
        self.map = flask_app.url_map.bind('')
        self.prefix = prefix
        self.blacklist = blacklist

    def _metric_name(self, path, method):
        path = path.split('?')[0]
        match = self.map.match(path, method)
        return '{}{}'.format(self.prefix and '{}.'.format(self.prefix) or '', match[0])

    def _is_excluded(self, path, method):
        match = self.map.match(path, method)
        return match[0] in self.blacklist

    def __call__(self, environ, start_response):
        metric_name = None
        blacklisted = False
        try:
            metric_name = self._metric_name(environ['PATH_INFO'],
                                            environ['REQUEST_METHOD'])
            blacklisted = self._is_excluded(environ['PATH_INFO'],
                                            environ['REQUEST_METHOD'])
        except Exception:
            pass

        start = time.time()
        try:
            iterable = self.wsgi_app(environ, start_response)
            for event in iterable:
                # make sure to yield each element in the iterable so
                # we're only emitting a timer once all of the request has
                # been sent to the client
                yield event
            end = time.time()
        finally:
            # http://blog.dscpl.com.au/2012/10/obligations-for-calling-close-on.html
            if hasattr(iterable, 'close') and callable(iterable.close):
                iterable.close()

        if metric_name and not blacklisted:
            elapsed = int((end - start) * 1000.0)
            self.statsd.timing(metric_name, elapsed)