majuna/app/__init__.py

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)

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"))


if __name__ == '__main__':
    app.run()