Flask extension that collects metrics using Metrology and push them to Graphite

flask_metrology.py

#NOTE This is alpha quality code. Works only on my workstation.
#NOTE Some details looks weird becase code is ported from my FlaskScales extension with minimal changes
from __future__ import unicode_literals
import os
import socket
from urlparse import urlparse
import logging
from flask import _request_ctx_stack
from flask.signals import got_request_exception, request_started, request_finished, template_rendered
from astrolabe import Interval
from metrology import Metrology
from metrology.reporter import LoggerReporter, GraphiteReporter
from common.ps_utils import get_proc_name


class FlaskMetrology(object):
    def init_app(self, app):
        """
        @param flask.app.Flask app: Instance of Flask application to decorate
        """
        self.app = app
        self.process_name = get_proc_name()
        self.graphite_url = app.config.get('METROLOGY_GRAPHITE_URL') \
            or os.environ.get('METROLOGY_GRAPHITE_URL')
        #TODO implement METROLOGY_GRAPHITE_ALLOW (and maybe METROLOGY_GRAPHITE_DENY) settings
        self.graphite_push_period = app.config.get('METROLOGY_GRAPHITE_PUSH_PERIOD') \
            or os.environ.get('METROLOGY_GRAPHITE_PUSH_PERIOD')
        self.graphite_prefix = app.config.get('METROLOGY_GRAPHITE_PREFIX') \
            or os.environ.get('METROLOGY_GRAPHITE_PREFIX') \
            or 'metrology.{}.{}'.format(socket.getfqdn().lower().replace('.', '_'), self.process_name)
        self.server_name = app.config.get('METROLOGY_SERVER_NAME') \
            or os.environ.get('METROLOGY_SERVER_NAME') or socket.getfqdn()
        #TODO develop satus page which will render json with current metrics
        self.url_prefix = app.config.get('METROLOGY_UI_PREFIX') \
            or os.environ.get('METROLOGY_UI_PREFIX') or '/status'
        self.stats_path = app.config.get('METROLOGY_STATS_PATH') \
            or os.environ.get('METROLOGY_STATS_PATH') or '/flask'

        Metrology.stop()
        app.before_first_request(self.handle_before_first_request)
        #TODO Use METROLOGY_STATS_PATH variable instead of hardcoded prefix
        self.request_2xx_timer = Metrology.timer('flask.http_code_2xx')
        self.request_3xx_timer = Metrology.timer('flask.http_code_3xx')
        self.request_4xx_timer = Metrology.timer('flask.http_code_4xx')
        self.request_5xx_timer = Metrology.timer('flask.http_code_5xx')
        self.request_per_sec = Metrology.meter('request_per_sec')

        #TODO drop this experimental code
        self.request_all_timer = Metrology.timer('request_time')
        self.request_all_utimer = Metrology.utilization_timer('request_time_u')

        template_rendered.connect(self.handle_template_rendered, sender=app)
        request_started.connect(self.handle_request_started, sender=app)
        request_finished.connect(self.handle_request_finished, sender=app)
        got_request_exception.connect(self.handle_request_error, sender=app)
        # TODO fire signals when calling db/redis/http
        # http://flask.pocoo.org/docs/signals/#creating-signals
        # http://flask.pocoo.org/docs/api/#flask.signals.Namespace

        #self.logger = LoggerReporter(logger=self.app.logger, level=logging.INFO, interval=10)
        #self.logger.start()

        if self.graphite_url:
            url = urlparse(self.graphite_url)
            self.graphite = GraphiteReporter(
                host=url.hostname, port=url.port or 2003, prefix=self.graphite_prefix,
                interval=self.graphite_push_period, pickle=False)
            self.graphite.start()

        # Flask app must have link to extensions due to subscription weak reference
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['metrology'] = self

    def handle_before_first_request(self):
        self.endpoint_timers = dict()
        endpoints = (rule.endpoint for rule in self.app.url_map.iter_rules())
        for endpoint in endpoints:
            timer = 'flask.endpoints.{}'.format(endpoint.replace('.', '__'))
            self.endpoint_timers[endpoint] = Metrology.timer(timer)

    def handle_template_rendered(self, sender, **kwargs):
        #TODO consider drop this handler
        pass

    def handle_request_started(self, sender, **kwargs):
        ctx = _request_ctx_stack.top
        ctx.metrology_request_time_interval = Interval.now()

    def handle_request_finished(self, sender, **kwargs):
        ctx = _request_ctx_stack.top
        duration = ctx.metrology_request_time_interval.stop()
        http_code = kwargs['response'].status_code
        if 200 <= http_code < 300:
            self.request_2xx_timer.update(duration)
        elif 300 <= http_code < 400:
            self.request_3xx_timer.update(duration)
        elif 400 <= http_code < 500:
            self.request_4xx_timer.update(duration)
        elif 500 <= http_code < 600:
            #TODO Check is this code will be called when user-code return HTTP 500 response w/o exception
            self.request_5xx_timer.update(duration)
        #TODO consider split each endpoint timer into 2xx, 3xx, 4xx, 5xx timers
        if ctx.request.url_rule:
            endpoint = ctx.request.url_rule.endpoint
            self.endpoint_timers[endpoint].update(duration)
        self.request_all_timer.update(duration)
        self.request_all_utimer.update(duration)

    def handle_request_error(self, sender, **kwargs):
        ctx = _request_ctx_stack.top
        duration = ctx.metrology_request_time_interval.stop()
        self.request_5xx_timer.update(duration)
        #exc_info = kwargs.get('exc_info') or kwargs.get('exception')
        #data = get_data_from_request(request)
        #TODO track hits by error type name or full dotted type name
        #TODO investigate how to track validation errors