dream

1
from flask import request, Response
2
from prometheus_client import generate_latest, Counter, Histogram, Gauge, CONTENT_TYPE_LATEST
3
import time
4

5
REQUEST_COUNT = Counter(
6
    "http_request_count",
7
    "App Request Count",
8
    [
9
        "method",
10
        "endpoint",
11
        "code",
12
    ],
13
)
14

15
REQUEST_LATENCY = Histogram("http_request_latency_seconds", "Request latency", ["endpoint"])
16

17
REQUEST_IN_PROGRESS = Gauge(
18
    "http_request_in_progress",
19
    "Requests in progress",
20
    [
21
        "endpoint",
22
    ],
23
)
24

25

26
def do_not_track(func):
27
    func._do_not_track = True
28
    return func
29

30

31
def setup_metrics(app):
32
    # readiness endpoint
33
    @app.route("/ready", methods=["GET"])
34
    @do_not_track
35
    def ready():
36
        return Response("OK")
37

38
    # liveness endpoint
39
    @app.route("/health", methods=["GET"])
40
    @do_not_track
41
    def health():
42
        return Response("OK")
43

44
    # metrics endpoint
45
    @app.route("/metrics", methods=["GET"])
46
    @do_not_track
47
    def metrics():
48
        return Response(generate_latest(), mimetype=CONTENT_TYPE_LATEST)
49

50
    def ignore_endpoint():
51
        view_func = app.view_functions.get(request.endpoint)
52
        return hasattr(view_func, "_do_not_track")
53

54
    def before_request():
55
        if ignore_endpoint():
56
            return
57
        request.start_time = time.time()
58
        REQUEST_IN_PROGRESS.labels(request.path).inc()
59

60
    def after_request(response):
61
        if ignore_endpoint():
62
            return response
63
        resp_time = time.time() - request.start_time
64
        REQUEST_COUNT.labels(request.method, request.path, response.status_code).inc()
65
        REQUEST_LATENCY.labels(request.path).observe(resp_time)
66
        REQUEST_IN_PROGRESS.labels(request.path).dec()
67
        return response
68

69
    app.before_request(before_request)
70
    app.after_request(after_request)
71