diff --git a/app/__init__.py b/app/__init__.py index 1eb591e..5f1a446 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -30,9 +30,9 @@ app.register_blueprint(tfstate, url_prefix="/tfstate") if 'nose' not in sys.modules.keys() and sys.argv[1] != "db": from app.metrics import DefinedProxiesCollector, BlockedProxiesCollector, AutomationCollector - REGISTRY.register(DefinedProxiesCollector()) # type: ignore[no-untyped-call] - REGISTRY.register(BlockedProxiesCollector()) # type: ignore[no-untyped-call] - REGISTRY.register(AutomationCollector()) # type: ignore[no-untyped-call] + REGISTRY.register(DefinedProxiesCollector(app)) # type: ignore[no-untyped-call] + REGISTRY.register(BlockedProxiesCollector(app)) # type: ignore[no-untyped-call] + REGISTRY.register(AutomationCollector(app)) # type: ignore[no-untyped-call] @app.route('/') diff --git a/app/metrics.py b/app/metrics.py index 3a2b2b3..a4c9d72 100644 --- a/app/metrics.py +++ b/app/metrics.py @@ -1,16 +1,24 @@ -from typing import Iterator +from typing import Iterator, Optional +from flask import Flask from prometheus_client.metrics_core import GaugeMetricFamily, CounterMetricFamily, Metric from sqlalchemy import text -from app import app from app.extensions import db from app.models.automation import Automation, AutomationState +class FlaskCollector: + def __init__(self, app=None): + self.app: Optional[Flask] = None + + def init_app(self, app): + self.app = app + + class DefinedProxiesCollector: def collect(self) -> Iterator[Metric]: - with app.app_context(): + with self.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, @@ -32,7 +40,7 @@ class DefinedProxiesCollector: class BlockedProxiesCollector: def collect(self) -> Iterator[Metric]: - with app.app_context(): + with self.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, @@ -56,7 +64,7 @@ class BlockedProxiesCollector: class AutomationCollector: def collect(self) -> Iterator[Metric]: - with app.app_context(): + with self.app.app_context(): c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)", labels=['automation_name']) # type: ignore[no-untyped-call] automations = Automation.query.all()