Refactor migrations as custom collector

django_prometheus/apps.py

@@ -2,8 +2,9 @@
 from django.conf import settings
 
 import django_prometheus
+from prometheus_client.core import REGISTRY
 from django_prometheus.exports import SetupPrometheusExportsFromConfig
-from django_prometheus.migrations import ExportMigrations
+from django_prometheus.migrations import MigrationCollector
 
 
 class DjangoPrometheusConfig(AppConfig):
@@ -21,4 +22,4 @@ def ready(self):
         """
         SetupPrometheusExportsFromConfig()
         if getattr(settings, "PROMETHEUS_EXPORT_MIGRATIONS", False):
-            ExportMigrations()
+            REGISTRY.register(MigrationCollector())

django_prometheus/migrations.py

@@ -1,51 +1,39 @@
 from django.db import connections
 from django.db.backends.dummy.base import DatabaseWrapper
-from prometheus_client import Gauge
-
-from django_prometheus.conf import NAMESPACE
-
-unapplied_migrations = Gauge(
-    "django_migrations_unapplied_total",
-    "Count of unapplied migrations by database connection",
-    ["connection"],
-    namespace=NAMESPACE,
-)
-
-applied_migrations = Gauge(
-    "django_migrations_applied_total",
-    "Count of applied migrations by database connection",
-    ["connection"],
-    namespace=NAMESPACE,
-)
-
-
-def ExportMigrationsForDatabase(alias, executor):
-    plan = executor.migration_plan(executor.loader.graph.leaf_nodes())
-    unapplied_migrations.labels(alias).set(len(plan))
-    applied_migrations.labels(alias).set(len(executor.loader.applied_migrations))
-
-
-def ExportMigrations():
-    """Exports counts of unapplied migrations.
-
-    This is meant to be called during app startup, ideally by
-    django_prometheus.apps.AppConfig.
-    """
-
-    # Import MigrationExecutor lazily. MigrationExecutor checks at
-    # import time that the apps are ready, and they are not when
-    # django_prometheus is imported. ExportMigrations() should be
-    # called in AppConfig.ready(), which signals that all apps are
-    # ready.
-    from django.db.migrations.executor import MigrationExecutor
-
-    if "default" in connections and (isinstance(connections["default"], DatabaseWrapper)):
+from prometheus_client.core import GaugeMetricFamily, REGISTRY
+from prometheus_client.registry import Collector
+
+class MigrationCollector(Collector):
+    def collect(self):
+        from django.db.migrations.executor import MigrationExecutor
+
+        applied_migrations = GaugeMetricFamily(
+            "django_migrations_applied_total",
+            "Count of applied migrations by database connection",
+            labels=["connection"]
+        )
+
+        unapplied_migrations = GaugeMetricFamily(
+            "django_migrations_unapplied_total",
+            "Count of unapplied migrations by database connection",
+            labels=["connection"]
+        )
+
+        if "default" in connections and isinstance(connections["default"], DatabaseWrapper):
         # This is the case where DATABASES = {} in the configuration,
         # i.e. the user is not using any databases. Django "helpfully"
         # adds a dummy database and then throws when you try to
         # actually use it. So we don't do anything, because trying to
         # export stats would crash the app on startup.
-        return
-    for alias in connections.databases:
-        executor = MigrationExecutor(connections[alias])
-        ExportMigrationsForDatabase(alias, executor)
+            return
+
+        for alias in connections.databases:
+            executor = MigrationExecutor(connections[alias])
+            applied_migrations_count = len(executor.loader.applied_migrations)
+            unapplied_migrations_count = len(executor.migration_plan(executor.loader.graph.leaf_nodes()))
+
+            applied_migrations.add_metric([alias], applied_migrations_count)
+            unapplied_migrations.add_metric([alias], unapplied_migrations_count)
+
+        yield applied_migrations
+        yield unapplied_migrations

django_prometheus/tests/end2end/testapp/test_migrations.py

@@ -1,9 +1,10 @@
-from unittest.mock import MagicMock
+from unittest.mock import patch
 
 import pytest
+from prometheus_client import CollectorRegistry
+
+from django_prometheus.migrations import MigrationCollector
 
-from django_prometheus.migrations import ExportMigrationsForDatabase
-from django_prometheus.testutils import assert_metric_equal
 
 
 def M(metric_name):
@@ -19,19 +20,38 @@ def M(metric_name):
 class TestMigrations:
     """Test migration counters."""
 
-    def test_counters(self):
-        executor = MagicMock()
-        executor.migration_plan = MagicMock()
-        executor.migration_plan.return_value = set()
-        executor.loader.applied_migrations = {"a", "b", "c"}
-        ExportMigrationsForDatabase("fakedb1", executor)
-        assert executor.migration_plan.call_count == 1
-        executor.migration_plan = MagicMock()
-        executor.migration_plan.return_value = {"a"}
-        executor.loader.applied_migrations = {"b", "c"}
-        ExportMigrationsForDatabase("fakedb2", executor)
-
-        assert_metric_equal(3, M("applied_total"), connection="fakedb1")
-        assert_metric_equal(0, M("unapplied_total"), connection="fakedb1")
-        assert_metric_equal(2, M("applied_total"), connection="fakedb2")
-        assert_metric_equal(1, M("unapplied_total"), connection="fakedb2")
+    @patch('django.db.migrations.executor.MigrationExecutor')
+    def test_counters(self, MockMigrationExecutor):
+
+        mock_executor = MockMigrationExecutor.return_value
+        mock_executor.migration_plan.return_value = set()
+        mock_executor.loader.applied_migrations = {"a", "b", "c"}
+
+        test_registry = CollectorRegistry()
+        collector = MigrationCollector()
+        test_registry.register(collector)
+
+        metrics = list(collector.collect())
+
+        applied_metric = next((m for m in metrics if m.name == M("applied_total")), None)
+        unapplied_metric = next((m for m in metrics if m.name == M("unapplied_total")), None)
+
+        assert applied_metric.samples[0].value == 3
+        assert applied_metric.samples[0].labels == {"connection": "default"}
+
+        assert unapplied_metric.samples[0].value == 0
+        assert unapplied_metric.samples[0].labels == {"connection": "default"}
+
+        mock_executor.migration_plan.return_value = {"a"}
+        mock_executor.loader.applied_migrations = {"b", "c"}
+
+        metrics = list(collector.collect())
+
+        applied_metric = next((m for m in metrics if m.name == M("applied_total")), None)
+        unapplied_metric = next((m for m in metrics if m.name == M("unapplied_total")), None)
+
+        assert applied_metric.samples[0].value == 2
+        assert applied_metric.samples[0].labels == {"connection": "default"}
+
+        assert unapplied_metric.samples[0].value == 1
+        assert unapplied_metric.samples[0].labels == {"connection": "default"}