majuna/app/__init__.py

import os
import sys
from typing import Iterator

import yaml
from flask import Flask, redirect, send_from_directory, url_for
from flask.typing import ResponseReturnValue
from prometheus_client import CollectorRegistry, Metric, make_wsgi_app
from prometheus_client.metrics_core import CounterMetricFamily, GaugeMetricFamily
from prometheus_client.registry import REGISTRY, Collector
from prometheus_flask_exporter import PrometheusMetrics
from sqlalchemy import text
from sqlalchemy.exc import SQLAlchemyError
from werkzeug.middleware.dispatcher import DispatcherMiddleware

from app.api import api
from app.extensions import bootstrap, db, migrate
from app.models.automation import Automation, AutomationState
from app.portal import portal
from app.portal.report import report
from app.tfstate import tfstate

app = Flask(__name__)
app.config.from_file("../config.yaml", load=yaml.safe_load)

# create new registry to avoid multiple metrics registration in global REGISTRY
registry = CollectorRegistry()
metrics = PrometheusMetrics(app, registry=registry)

app.wsgi_app = DispatcherMiddleware(  # type: ignore[method-assign]
    app.wsgi_app, {"/metrics": make_wsgi_app(registry)}
)

# register default collectors to our new registry
collectors = list(REGISTRY._collector_to_names.keys())
for collector in collectors:
    registry.register(collector)

db.init_app(app)
migrate.init_app(app, db, render_as_batch=True)
bootstrap.init_app(app)

app.register_blueprint(api, url_prefix="/api")
app.register_blueprint(portal, url_prefix="/portal")
app.register_blueprint(tfstate, url_prefix="/tfstate")
app.register_blueprint(report, url_prefix="/report")


def not_migrating() -> bool:
    return len(sys.argv) < 2 or sys.argv[1] != "db"


class DefinedProxiesCollector(Collector):
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            ok = GaugeMetricFamily(
                "database_collector",
                "Status of a database collector (0: bad, 1: good)",
                labels=["collector"],
            )
            try:
                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,
                                   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
                    ok.add_metric(["defined_proxies"], 1)
            except SQLAlchemyError:
                ok.add_metric(["defined_proxies"], 0)
            yield ok


class BlockedProxiesCollector(Collector):
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            ok = GaugeMetricFamily(
                "database_collector",
                "Status of a database collector (0: bad, 1: good)",
                labels=["collector"],
            )
            try:
                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
                    ok.add_metric(["deprecated_proxies"], 0)
            except SQLAlchemyError:
                ok.add_metric(["deprecated_proxies"], 0)
            yield ok


class AutomationCollector(Collector):
    def collect(self) -> Iterator[Metric]:
        with app.app_context():
            ok = GaugeMetricFamily(
                "database_collector",
                "Status of a database collector (0: bad, 1: good)",
                labels=["collector"],
            )
            try:
                state = GaugeMetricFamily(
                    "automation_state",
                    "The automation state (0: idle, 1: running, 2: error)",
                    labels=["automation_name"],
                )
                enabled = GaugeMetricFamily(
                    "automation_enabled",
                    "Whether an automation is enabled (0: disabled, 1: enabled)",
                    labels=["automation_name"],
                )
                next_run = GaugeMetricFamily(
                    "automation_next_run",
                    "The timestamp of the next run of the automation",
                    labels=["automation_name"],
                )
                last_run_start = GaugeMetricFamily(
                    "automation_last_run_start",
                    "The timestamp of the last run of the automation ",
                    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:
                        state.add_metric([automation.short_name], 0)
                    elif automation.state == AutomationState.RUNNING:
                        state.add_metric([automation.short_name], 1)
                    else:
                        state.add_metric([automation.short_name], 2)
                    enabled.add_metric(
                        [automation.short_name], 1 if automation.enabled else 0
                    )
                    if automation.next_run:
                        next_run.add_metric(
                            [automation.short_name], automation.next_run.timestamp()
                        )
                    else:
                        next_run.add_metric([automation.short_name], 0)
                    if automation.last_run:
                        last_run_start.add_metric(
                            [automation.short_name], automation.last_run.timestamp()
                        )
                    else:
                        last_run_start.add_metric([automation.short_name], 0)
                yield state
                yield enabled
                yield next_run
                yield last_run_start
                ok.add_metric(["automation_state"], 1)
            except SQLAlchemyError:
                ok.add_metric(["automation_state"], 0)
            yield ok


# register all custom collectors to registry
if not_migrating() and "DISABLE_METRICS" not in os.environ:
    registry.register(DefinedProxiesCollector())
    registry.register(BlockedProxiesCollector())
    registry.register(AutomationCollector())


@app.route("/ui")
def redirect_ui() -> ResponseReturnValue:
    return redirect("/ui/")


@app.route("/ui/", defaults={"path": ""})
@app.route("/ui/<path:path>")
def serve_ui(path: str) -> ResponseReturnValue:
    if path != "" and os.path.exists("app/static/ui/" + path):
        return send_from_directory("static/ui", path)
    else:
        return send_from_directory("static/ui", "index.html")


@app.route("/")
def index() -> ResponseReturnValue:
    # TODO: update to point at new UI when ready
    return redirect(url_for("portal.portal_home"))


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