Add prometheus metrics to flask app

matl_online/app.py

@@ -7,7 +7,7 @@
 from matl_online import public
 from matl_online.assets import assets
 from matl_online.commands import register_commands
-from matl_online.extensions import celery, csrf, db, migrate, rollbar, socketio
+from matl_online.extensions import celery, csrf, db, migrate, rollbar, socketio, metrics
 from matl_online.settings import Config, get_celery_configuration, get_config
 
 
@@ -40,6 +40,9 @@ def register_extensions(app: Flask) -> None:
     celery.conf.update(get_celery_configuration(app.config))
     csrf.init_app(app)
 
+    metrics.init_app(app)
+    metrics.info("app_info", "Application info", version=app.config.get("APP_VERSION"))
+
 
 def register_rollbar(app: Flask) -> None:
     rollbar.init(
@@ -56,3 +59,6 @@ def register_blueprints(app: Flask) -> None:
     """Register Flask blueprints."""
     app.register_blueprint(public.views.blueprint)
     return None
+
+
+__all__ = ["metrics"]

matl_online/extensions.py

@@ -6,11 +6,21 @@
 from flask_socketio import SocketIO  # type: ignore[import]
 from flask_sqlalchemy import SQLAlchemy
 from flask_wtf import CSRFProtect  # type: ignore[import]
+from prometheus_flask_exporter import PrometheusMetrics  # type: ignore[import]
 
 db = SQLAlchemy()
 migrate = Migrate()
 socketio = SocketIO()
 celery = Celery(__name__)
 csrf = CSRFProtect()
+metrics = PrometheusMetrics.for_app_factory()
 
-__all__ = ["rollbar", "db", "migrate", "socketio", "celery", "csrf"]
+__all__ = [
+    "celery",
+    "csrf",
+    "db",
+    "metrics",
+    "migrate",
+    "rollbar",
+    "socketio",
+]

matl_online/public/views.py

@@ -17,7 +17,7 @@
 from wtforms import ValidationError  # type: ignore
 
 from matl_online.errors import InvalidVersion
-from matl_online.extensions import celery, csrf, socketio
+from matl_online.extensions import celery, csrf, socketio, metrics
 from matl_online.matl.documentation import help_file
 from matl_online.matl.releases import refresh_releases
 from matl_online.public.models import Release
@@ -180,13 +180,15 @@ def share() -> Tuple[Response, int]:
 
 
 @socketio.on("connect")  # type: ignore[misc]
+@metrics.counter("socketio_connections", "SocketIO Events")  # type: ignore[misc]
 def connected() -> None:
     """Send an event to the client with the ID of their session."""
     session_id = rooms()[0]
     emit("connection", {"session_id": session_id})
 
 
 @socketio.on("kill")  # type: ignore[misc]
+@metrics.counter("socketio_kill_events", "SocketIO Kill Events")  # type: ignore[misc]
 def kill_task(data: Any) -> None:
     """Triggered when a kill message is sent to kill a task."""
     taskid = session.get("taskid", None)
@@ -201,6 +203,7 @@ def kill_task(data: Any) -> None:
 
 
 @socketio.on("submit")  # type: ignore[misc]
+@metrics.counter("socketio_submit_events", "SocketIO Submit Events")  # type: ignore[misc]
 def submit_job(data: Dict[str, Any]) -> None:
     """Submit some code and inputs for interpretation."""
     # If we already have a task disable submitting

matl_online/settings.py

@@ -21,6 +21,8 @@ class Config(object):
 
     ENV: str = "NONE"
 
+    APP_VERSION = os.environ.get("APP_VERSION", "")
+
     IMGUR_CLIENT_ID = os.environ.get("MATL_ONLINE_IMGUR_CLIENT_ID")
 
     SECRET_KEY = str(uuid.uuid4())

requirements/prod.txt

@@ -49,3 +49,6 @@ pygithub==1.57
 # Error Tracking
 blinker==1.5
 rollbar==0.16.3
+
+# Metrics
+prometheus-flask-exporter==0.22.0

tests/conftest.py

@@ -6,15 +6,19 @@
 from typing import Any, Generator
 from unittest.mock import Mock
 
+import prometheus_client
 import pytest
 from flask import Flask
 from flask_socketio import SocketIOTestClient  # type: ignore
 from flask_sqlalchemy import SQLAlchemy
+from prometheus_flask_exporter import PrometheusMetrics  # type: ignore[import]
 from pytest_mock.plugin import MockerFixture
 from webtest import TestApp  # type: ignore
 
 os.environ["MATL_ONLINE_ENV"] = "test"
 
+from matl_online import app as matl_online_app  # noqa: E402
+from matl_online import extensions  # noqa: E402
 from matl_online.app import create_app  # noqa: E402
 from matl_online.database import db as _db  # noqa: E402
 from matl_online.extensions import socketio  # noqa: E402
@@ -35,8 +39,16 @@ def socketio_client(app: Flask) -> SocketIOTestClient:
 
 
 @pytest.fixture(scope="function")
-def app() -> Generator[Flask, None, None]:
+def metrics() -> PrometheusMetrics:
+    registry = prometheus_client.CollectorRegistry(auto_describe=True)
+    return PrometheusMetrics.for_app_factory(registry=registry)
+
+
+@pytest.fixture(scope="function")
+def app(metrics: PrometheusMetrics) -> Generator[Flask, None, None]:
     """Flask app instance."""
+    matl_online_app.metrics = metrics
+    extensions.metrics = metrics
     app_instance = create_app(TestConfig)
     context = app_instance.test_request_context()
     context.push()  # type: ignore[attr-defined]

tests/test_config.py

@@ -2,29 +2,32 @@
 
 import os
 
-from matl_online.app import create_app
+from matl_online import app
 from matl_online.settings import (
     DevConfig,
     ProdConfig,
     _get_cors_allowed_origins,
     get_config,
 )
+from prometheus_flask_exporter import PrometheusMetrics  # type: ignore[import]
 
 
-def test_production_config() -> None:
+def test_production_config(metrics: PrometheusMetrics) -> None:
     """Production config."""
-    app = create_app(ProdConfig)
-    assert app.config["ENV"] == "prod"
-    assert app.config["DEBUG"] is False
-    assert app.config["ASSETS_DEBUG"] is False
+    app.metrics = metrics
+    flask_app = app.create_app(ProdConfig)
+    assert flask_app.config["ENV"] == "prod"
+    assert flask_app.config["DEBUG"] is False
+    assert flask_app.config["ASSETS_DEBUG"] is False
 
 
-def test_dev_config() -> None:
+def test_dev_config(metrics: PrometheusMetrics) -> None:
     """Development config."""
-    app = create_app(DevConfig)
-    assert app.config["ENV"] == "dev"
-    assert app.config["DEBUG"] is True
-    assert app.config["ASSETS_DEBUG"] is True
+    app.metrics = metrics
+    flask_app = app.create_app(DevConfig)
+    assert flask_app.config["ENV"] == "dev"
+    assert flask_app.config["DEBUG"] is True
+    assert flask_app.config["ASSETS_DEBUG"] is True
 
 
 def test_prod_config_lookup() -> None: