metrics: merge back to __init__.py to avoid context issues

app/__init__.py

@@ -1,17 +1,22 @@
+from typing import Iterator
+
 from flask import Flask, redirect, url_for
 from flask.typing import ResponseReturnValue
+from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily
+from prometheus_client.registry import Collector
+from sqlalchemy import text
 from werkzeug.middleware.dispatcher import DispatcherMiddleware
-from prometheus_client import make_wsgi_app, REGISTRY
+from prometheus_client import make_wsgi_app, REGISTRY, Metric
 import yaml
 import sys
 
 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)
 
@@ -31,12 +36,73 @@ def not_migrating() -> bool:
     return len(sys.argv) < 2 or sys.argv[1] != "db"
 
 
-if 'nose' not in sys.modules.keys() and not_migrating():
-    from app.metrics import DefinedProxiesCollector, BlockedProxiesCollector, AutomationCollector
+class DefinedProxiesCollector(Collector):
+    def collect(self) -> Iterator[Metric]:
+        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
 
-    REGISTRY.register(DefinedProxiesCollector(app))
-    REGISTRY.register(BlockedProxiesCollector(app))
-    REGISTRY.register(AutomationCollector(app))
+
+class BlockedProxiesCollector(Collector):
+    def collect(self) -> Iterator[Metric]:
+        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("deprecated_proxies",
+                                        "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(Collector):
+    def collect(self) -> Iterator[Metric]:
+        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('/')