
Each worker now uses its own registry instead of the global REGISTRY to avoid duplicate metric registration, since metrics are served from the database.
136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
import os
|
|
import sys
|
|
from typing import Iterator
|
|
|
|
import yaml
|
|
from flask import Flask, redirect, url_for, send_from_directory
|
|
from flask.typing import ResponseReturnValue
|
|
from prometheus_client import make_wsgi_app, Metric, CollectorRegistry
|
|
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 app.api import api
|
|
from app.extensions import bootstrap
|
|
from app.extensions import db
|
|
from app.extensions import 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)
|
|
|
|
app.wsgi_app = DispatcherMiddleware(app.wsgi_app, { # type: ignore[method-assign]
|
|
'/metrics': make_wsgi_app()
|
|
})
|
|
|
|
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():
|
|
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(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
|
|
|
|
|
|
if not_migrating() and 'DISABLE_METRICS' not in os.environ:
|
|
registry = CollectorRegistry()
|
|
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()
|