From 965d79b08fa611d86c3507c9304bac9eb89fab01 Mon Sep 17 00:00:00 2001
From: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date: Wed, 15 Jun 2022 14:52:22 +0200
Subject: [PATCH] add prometheus endpoint to analysis server

---
 pyinfra/server/server.py | 12 +++++++++++-
 requirements.txt         |  1 +
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/pyinfra/server/server.py b/pyinfra/server/server.py
index 21781e2..4519e5c 100644
--- a/pyinfra/server/server.py
+++ b/pyinfra/server/server.py
@@ -3,6 +3,7 @@ from typing import Dict
 
 from flask import Flask, jsonify, request
 from funcy import merge
+from prometheus_client import generate_latest, Summary
 
 from pyinfra.server.buffering.stream import FlatStreamBuffer
 from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction
@@ -54,6 +55,10 @@ def __set_up_processing_server_impl(operation2function: Dict[str, QueuedStreamFu
     pickup_operation2processor = {pickup_suffix(op): prc for op, prc in operation2processor.items()}
     operation2processor = merge(submit_operation2processor, pickup_operation2processor)
 
+    pickup_operation2metric = {
+        op: Summary(f"redactmanager_{op}_seconds", f"Time spent on {op}.") for op in operation2processor
+    }
+
     @app.route("/ready", methods=["GET"])
     def ready():
         resp = jsonify("OK")
@@ -66,12 +71,17 @@ def __set_up_processing_server_impl(operation2function: Dict[str, QueuedStreamFu
         resp.status_code = 200
         return resp
 
+    @app.route("/prometheus", methods=["GET"])
+    def prometheus():
+        return generate_latest()
+
     @app.route("/<operation>", methods=["POST", "PATCH"])
     def submit(operation):
         return operation2processor[operation].push(request)
 
     @app.route("/<operation>", methods=["GET"])
     def pickup(operation):
-        return operation2processor[operation].pop()
+        with pickup_operation2metric[operation.replace("_pickup", "")].time():
+            return operation2processor[operation].pop()
 
     return app
diff --git a/requirements.txt b/requirements.txt
index 68a8956..59682fe 100755
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,3 +17,4 @@ PyMuPDF==1.19.6
 more-itertools==8.12.0
 numpy==1.22.3
 Pillow==9.0.1
+prometheus-client==0.13.1