metrics: refactor app context into base class

This commit is contained in:
Iain Learmonth 2022-11-02 13:46:39 +00:00
parent 28c07b7d82
commit 0ff61721d9

View file

@ -1,3 +1,4 @@
from abc import abstractmethod
from typing import Iterator, Optional from typing import Iterator, Optional
from flask import Flask from flask import Flask
@ -9,16 +10,27 @@ from app.models.automation import Automation, AutomationState
class FlaskCollector: class FlaskCollector:
def __init__(self, app=None): _app: Optional[Flask]
self.app: Optional[Flask] = None
def init_app(self, app): def __init__(self, app: Optional[Flask] = None) -> None:
self.app = app self._app = app
def init_app(self, app: Flask) -> None:
self._app = app
class DefinedProxiesCollector:
def collect(self) -> Iterator[Metric]: def collect(self) -> Iterator[Metric]:
with self.app.app_context(): 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() conn = db.engine.connect()
result = conn.execute(text(""" result = conn.execute(text("""
SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name, SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
@ -38,9 +50,8 @@ class DefinedProxiesCollector:
yield c yield c
class BlockedProxiesCollector: class BlockedProxiesCollector(FlaskCollector):
def collect(self) -> Iterator[Metric]: def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
with self.app.app_context():
with db.engine.connect() as conn: with db.engine.connect() as conn:
result = conn.execute(text(""" result = conn.execute(text("""
SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name, SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
@ -62,9 +73,8 @@ class BlockedProxiesCollector:
yield c yield c
class AutomationCollector: class AutomationCollector(FlaskCollector):
def collect(self) -> Iterator[Metric]: def collect_in_ctx(self, app: Flask) -> Iterator[Metric]:
with self.app.app_context():
c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)", c = GaugeMetricFamily("automation_state", "The automation state (0: idle, 1: running, 2: error)",
labels=['automation_name']) # type: ignore[no-untyped-call] labels=['automation_name']) # type: ignore[no-untyped-call]
automations = Automation.query.all() automations = Automation.query.all()