diff --git a/app/__init__.py b/app/__init__.py index 9cc211d..0b76ab1 100644 --- a/app/__init__.py +++ b/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('/') diff --git a/app/metrics.py b/app/metrics.py deleted file mode 100644 index ced81d2..0000000 --- a/app/metrics.py +++ /dev/null @@ -1,91 +0,0 @@ -from abc import abstractmethod -from typing import Iterator, Optional - -from flask import Flask -from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily, Metric -from prometheus_client.registry import Collector -from sqlalchemy import text - -from app.extensions import db -from app.models.automation import Automation, AutomationState - - -class FlaskCollector(Collector): - _app: Optional[Flask] - - def __init__(self, app: Optional[Flask] = None) -> None: - self._app = app - - def init_app(self, app: Flask) -> None: - self._app = app - - def collect(self) -> Iterator[Metric]: - if self._app: - with self._app.app_context(): - return self.collect_in_ctx(self._app) - raise RuntimeError("Flask collector was not initialised with app") - - @abstractmethod - def collect_in_ctx(self, app: Flask) -> Iterator[Metric]: - raise NotImplementedError() - - -class DefinedProxiesCollector(FlaskCollector): - def collect_in_ctx(self, app: Flask) -> Iterator[Metric]: - 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(FlaskCollector): - def collect_in_ctx(self, app: Flask) -> Iterator[Metric]: - 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(FlaskCollector): - def collect_in_ctx(self, app: Flask) -> Iterator[Metric]: - 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