From 6c25d8a86879699a933152ffe352c358323a50e9 Mon Sep 17 00:00:00 2001 From: luca-medeiros-reef Date: Wed, 19 Jun 2024 18:00:30 +0900 Subject: [PATCH 1/3] refactor migrations as custom collector --- django_prometheus/apps.py | 5 ++- django_prometheus/migrations.py | 76 ++++++++++++++------------------- 2 files changed, 35 insertions(+), 46 deletions(-) diff --git a/django_prometheus/apps.py b/django_prometheus/apps.py index 6a7ce8c9..4bc053d9 100644 --- a/django_prometheus/apps.py +++ b/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()) diff --git a/django_prometheus/migrations.py b/django_prometheus/migrations.py index 97908751..5e21e1ee 100644 --- a/django_prometheus/migrations.py +++ b/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 From ff348fb5ad9ae6cf9802df93cb3572be2acb21f5 Mon Sep 17 00:00:00 2001 From: luca-medeiros-reef Date: Wed, 19 Jun 2024 18:01:07 +0900 Subject: [PATCH 2/3] update test code --- .../tests/end2end/testapp/test_migrations.py | 58 +++++++++++++------ 1 file changed, 39 insertions(+), 19 deletions(-) diff --git a/django_prometheus/tests/end2end/testapp/test_migrations.py b/django_prometheus/tests/end2end/testapp/test_migrations.py index 6f5f1217..de675188 100644 --- a/django_prometheus/tests/end2end/testapp/test_migrations.py +++ b/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[1].value == 2 + assert applied_metric.samples[1].labels == {"connection": "test_db_1"} + + assert unapplied_metric.samples[1].value == 1 + assert unapplied_metric.samples[1].labels == {"connection": "test_db_1"} From a6007336c7403f83ade3ab0e1d926739afd302e1 Mon Sep 17 00:00:00 2001 From: luca-medeiros-reef Date: Wed, 19 Jun 2024 18:05:53 +0900 Subject: [PATCH 3/3] update connection name --- .../tests/end2end/testapp/test_migrations.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/django_prometheus/tests/end2end/testapp/test_migrations.py b/django_prometheus/tests/end2end/testapp/test_migrations.py index de675188..fca38073 100644 --- a/django_prometheus/tests/end2end/testapp/test_migrations.py +++ b/django_prometheus/tests/end2end/testapp/test_migrations.py @@ -50,8 +50,8 @@ def test_counters(self, MockMigrationExecutor): 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[1].value == 2 - assert applied_metric.samples[1].labels == {"connection": "test_db_1"} + assert applied_metric.samples[0].value == 2 + assert applied_metric.samples[0].labels == {"connection": "default"} - assert unapplied_metric.samples[1].value == 1 - assert unapplied_metric.samples[1].labels == {"connection": "test_db_1"} + assert unapplied_metric.samples[0].value == 1 + assert unapplied_metric.samples[0].labels == {"connection": "default"}