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 flask import Flask
@ -9,16 +10,27 @@ from app.models.automation import Automation, AutomationState
class FlaskCollector:
def __init__(self, app=None):
self.app: Optional[Flask] = None
_app: Optional[Flask]
def init_app(self, app):
self.app = app
def __init__(self, app: Optional[Flask] = None) -> None:
self._app = app
def init_app(self, app: Flask) -> None:
self._app = app
class DefinedProxiesCollector:
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()
result = conn.execute(text("""
SELECT origin.group_id, "group".group_name, proxy.provider, proxy.pool_id, pool.pool_name,
@ -38,9 +50,8 @@ class DefinedProxiesCollector:
yield c
class BlockedProxiesCollector:
def collect(self) -> Iterator[Metric]:
with self.app.app_context():
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,
@ -62,9 +73,8 @@ class BlockedProxiesCollector:
yield c
class AutomationCollector:
def collect(self) -> Iterator[Metric]:
with self.app.app_context():
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']) # type: ignore[no-untyped-call]
automations = Automation.query.all()