239 lines
9.3 KiB
Python
239 lines
9.3 KiB
Python
import os
|
|
import sys
|
|
from typing import Iterator
|
|
|
|
import yaml
|
|
from flask import Flask, redirect, send_from_directory, url_for
|
|
from flask.typing import ResponseReturnValue
|
|
from prometheus_client import CollectorRegistry, Metric, make_wsgi_app
|
|
from prometheus_client.metrics_core import CounterMetricFamily, GaugeMetricFamily
|
|
from prometheus_client.registry import REGISTRY, Collector
|
|
from prometheus_flask_exporter import PrometheusMetrics
|
|
from sqlalchemy import text
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from werkzeug.middleware.dispatcher import DispatcherMiddleware
|
|
|
|
from app.api import api
|
|
from app.extensions import bootstrap, db, 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)
|
|
|
|
# create new registry to avoid multiple metrics registration in global REGISTRY
|
|
registry = CollectorRegistry()
|
|
metrics = PrometheusMetrics(app, registry=registry)
|
|
|
|
app.wsgi_app = DispatcherMiddleware( # type: ignore[method-assign]
|
|
app.wsgi_app, {"/metrics": make_wsgi_app(registry)}
|
|
)
|
|
|
|
# register default collectors to our new registry
|
|
collectors = list(REGISTRY._collector_to_names.keys())
|
|
for collector in collectors:
|
|
registry.register(collector)
|
|
|
|
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():
|
|
ok = GaugeMetricFamily(
|
|
"database_collector",
|
|
"Status of a database collector (0: bad, 1: good)",
|
|
labels=["collector"],
|
|
)
|
|
try:
|
|
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,
|
|
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
|
|
ok.add_metric(["defined_proxies"], 1)
|
|
except SQLAlchemyError:
|
|
ok.add_metric(["defined_proxies"], 0)
|
|
yield ok
|
|
|
|
|
|
class BlockedProxiesCollector(Collector):
|
|
def collect(self) -> Iterator[Metric]:
|
|
with app.app_context():
|
|
ok = GaugeMetricFamily(
|
|
"database_collector",
|
|
"Status of a database collector (0: bad, 1: good)",
|
|
labels=["collector"],
|
|
)
|
|
try:
|
|
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
|
|
ok.add_metric(["deprecated_proxies"], 0)
|
|
except SQLAlchemyError:
|
|
ok.add_metric(["deprecated_proxies"], 0)
|
|
yield ok
|
|
|
|
|
|
class AutomationCollector(Collector):
|
|
def collect(self) -> Iterator[Metric]:
|
|
with app.app_context():
|
|
ok = GaugeMetricFamily(
|
|
"database_collector",
|
|
"Status of a database collector (0: bad, 1: good)",
|
|
labels=["collector"],
|
|
)
|
|
try:
|
|
state = GaugeMetricFamily(
|
|
"automation_state",
|
|
"The automation state (0: idle, 1: running, 2: error)",
|
|
labels=["automation_name"],
|
|
)
|
|
enabled = GaugeMetricFamily(
|
|
"automation_enabled",
|
|
"Whether an automation is enabled (0: disabled, 1: enabled)",
|
|
labels=["automation_name"],
|
|
)
|
|
next_run = GaugeMetricFamily(
|
|
"automation_next_run",
|
|
"The timestamp of the next run of the automation",
|
|
labels=["automation_name"],
|
|
)
|
|
last_run_start = GaugeMetricFamily(
|
|
"automation_last_run_start",
|
|
"The timestamp of the last run of the automation ",
|
|
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:
|
|
state.add_metric([automation.short_name], 0)
|
|
elif automation.state == AutomationState.RUNNING:
|
|
state.add_metric([automation.short_name], 1)
|
|
else:
|
|
state.add_metric([automation.short_name], 2)
|
|
enabled.add_metric(
|
|
[automation.short_name], 1 if automation.enabled else 0
|
|
)
|
|
if automation.next_run:
|
|
next_run.add_metric(
|
|
[automation.short_name], automation.next_run.timestamp()
|
|
)
|
|
else:
|
|
next_run.add_metric([automation.short_name], 0)
|
|
if automation.last_run:
|
|
last_run_start.add_metric(
|
|
[automation.short_name], automation.last_run.timestamp()
|
|
)
|
|
else:
|
|
last_run_start.add_metric([automation.short_name], 0)
|
|
yield state
|
|
yield enabled
|
|
yield next_run
|
|
yield last_run_start
|
|
ok.add_metric(["automation_state"], 1)
|
|
except SQLAlchemyError:
|
|
ok.add_metric(["automation_state"], 0)
|
|
yield ok
|
|
|
|
|
|
# register all custom collectors to registry
|
|
if not_migrating() and "DISABLE_METRICS" not in os.environ:
|
|
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()
|