From 6838596e8f6542c2e74aa27dc893fe0d4a0ac675 Mon Sep 17 00:00:00 2001
From: Iain Learmonth <irl@sr2.uk>
Date: Wed, 5 Oct 2022 15:25:51 +0100
Subject: [PATCH] metrics: very basic initial metrics

---
 app/__init__.py | 71 +++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/app/__init__.py b/app/__init__.py
index 0bb8b95..aefcbfd 100644
--- a/app/__init__.py
+++ b/app/__init__.py
@@ -1,15 +1,26 @@
 from flask import Flask, redirect, url_for
 from flask.typing import ResponseReturnValue
+from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily
+from sqlalchemy import text
+from werkzeug.middleware.dispatcher import DispatcherMiddleware
+from prometheus_client import make_wsgi_app, REGISTRY
 import yaml
 
 from app.extensions import db
 from app.extensions import migrate
 from app.extensions import bootstrap
+from app.models.automation import Automation, AutomationState
 from app.portal import portal
 from app.tfstate import tfstate
 
+
 app = Flask(__name__)
 app.config.from_file("../config.yaml", load=yaml.safe_load)
+
+app.wsgi_app = DispatcherMiddleware(app.wsgi_app, {
+    '/metrics': make_wsgi_app()
+})
+
 db.init_app(app)
 migrate.init_app(app, db, render_as_batch=True)
 bootstrap.init_app(app)
@@ -18,6 +29,66 @@ app.register_blueprint(portal, url_prefix="/portal")
 app.register_blueprint(tfstate, url_prefix="/tfstate")
 
 
+class DefinedProxiesCollector:
+    def collect(self):
+        with app.app_context():
+            conn = db.engine.connect()
+            result = conn.execute(text("""
+                    SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name, COUNT(proxy.id) FROM proxy, origin, pool, "group"
+                    WHERE proxy.origin_id = origin.id
+                        AND origin.group_id = "group".id
+                        AND proxy.pool_id = pool.id
+                        AND proxy.destroyed IS NULL
+                    GROUP BY origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name;
+                    """))
+            c = GaugeMetricFamily("defined_proxies", "Number of proxies currently defined for deployment", labels=['group_id', 'group_name', 'provider', 'pool_id', 'pool_name'])
+            for row in result:
+                c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4]], row[5])
+            yield c
+
+
+class BlockedProxiesCollector:
+    def collect(self):
+        with app.app_context():
+            with db.engine.connect() as conn:
+                result = conn.execute(text("""
+                        SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name, proxy.deprecation_reason, COUNT(proxy.id) FROM proxy, origin, pool, "group"
+                        WHERE proxy.origin_id = origin.id
+                            AND origin.group_id = "group".id
+                            AND proxy.pool_id = pool.id
+                            AND proxy.deprecated IS NOT NULL
+                GROUP BY origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name, proxy.deprecation_reason;
+                        """))
+                c = CounterMetricFamily(f"deprecated_proxies",
+                                      f"Number of proxies deprecated",
+                                      labels=['group_id', 'group_name', 'provider', 'pool_id', 'pool_name', 'deprecation_reason'])
+                for row in result:
+                    c.add_metric([str(row[0]), row[1], row[2], str(row[3]), row[4], row[5]], row[6])
+                yield c
+
+
+class AutomationCollector:
+    def collect(self):
+        with app.app_context():
+            c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)", labels=['automation_name'])
+            automations = Automation.query.all()
+            for automation in automations:
+                if automation.short_name in app.config['HIDDEN_AUTOMATIONS']:
+                    continue
+                if automation.state == AutomationState.IDLE:
+                    c.add_metric([automation.short_name], 0)
+                elif automation.state == AutomationState.RUNNING:
+                    c.add_metric([automation.short_name], 1)
+                else:
+                    c.add_metric([automation.short_name], 2)
+            yield c
+
+
+REGISTRY.register(DefinedProxiesCollector())
+REGISTRY.register(BlockedProxiesCollector())
+REGISTRY.register(AutomationCollector())
+
+
 @app.route('/')
 def index() -> ResponseReturnValue:
     return redirect(url_for("portal.portal_home"))