From 0bd4fa8099c9373b72247410d8b0d9320cd82c5b Mon Sep 17 00:00:00 2001
From: Nicolas Ledez <github.public@ledez.net>
Date: Mon, 2 Oct 2023 16:48:57 +0200
Subject: [PATCH] Add prometheus /metrics

---
 setup.cfg                         |  1 +
 zou/app/__init__.py               | 16 ++++++++++++++++
 zou/app/services/stats_service.py | 29 +++++++++++++++++++----------
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index 2e31c788d2..93d3022c7e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -58,6 +58,7 @@ install_requires =
     OpenTimelineIO==0.15.0
     orjson==3.9.7
     pillow==10.0.1
+    prometheus-flask-exporter==0.22.4
     psutil==5.9.5
     psycopg[binary]==3.1.12
     pygelf==0.4.2
diff --git a/zou/app/__init__.py b/zou/app/__init__.py
index 2dd2852471..9597bcc5a1 100644
--- a/zou/app/__init__.py
+++ b/zou/app/__init__.py
@@ -162,8 +162,24 @@ def add_permissions(_, payload):
 
 def load_api():
     from zou.app import api
+    from zou import __version__ as zou_version
 
     api.configure(app)
+    try:
+        from prometheus_flask_exporter.multiprocess import (
+            GunicornPrometheusMetrics,
+        )
+
+        metrics = GunicornPrometheusMetrics(
+            app, defaults_prefix="zou", group_by="url_rule"
+        )
+    except ValueError:
+        from prometheus_flask_exporter import RESTfulPrometheusMetrics
+
+        metrics = RESTfulPrometheusMetrics(
+            app, api, defaults_prefix="zou", group_by="url_rule"
+        )
+    metrics.info("zou_info", "Application info", version=zou_version)
 
     fs.mkdir_p(app.config["TMP_DIR"])
     configure_auth()
diff --git a/zou/app/services/stats_service.py b/zou/app/services/stats_service.py
index 51a417d2c7..0107258fd1 100644
--- a/zou/app/services/stats_service.py
+++ b/zou/app/services/stats_service.py
@@ -39,17 +39,26 @@
 
 
 def get_main_stats():
+    metric_number_of_video_previews = metrics.gauge(
+        "zou_number_of_video_previews", "Number of video previews"
+    )
+    number_of_video_previews = PreviewFile.query.filter(
+        PreviewFile.extension == "mp4"
+    ).count()
+    number_of_picture_previews = PreviewFile.query.filter(
+        PreviewFile.extension == "png"
+    ).count()
+    number_of_model_previews = PreviewFile.query.filter(
+        PreviewFile.extension == "obj"
+    ).count()
+    number_of_comments = Comment.query.count()
+
+    metric_number_of_video_previews.set(number_of_video_previews)
     return {
-        "number_of_video_previews": PreviewFile.query.filter(
-            PreviewFile.extension == "mp4"
-        ).count(),
-        "number_of_picture_previews": PreviewFile.query.filter(
-            PreviewFile.extension == "png"
-        ).count(),
-        "number_of_model_previews": PreviewFile.query.filter(
-            PreviewFile.extension == "obj"
-        ).count(),
-        "number_of_comments": Comment.query.count(),
+        "number_of_video_previews": number_of_video_previews,
+        "number_of_picture_previews": number_of_picture_previews,
+        "number_of_model_previews": number_of_model_previews,
+        "number_of_comments": number_of_comments,
     }