feat: add user agent to prom metrics

codecov/settings_base.py

@@ -49,7 +49,7 @@
 ]
 
 MIDDLEWARE = [
-    "django_prometheus.middleware.PrometheusBeforeMiddleware",
+    "core.middleware.AppMetricsBeforeMiddlewareWithUA",
     "django.middleware.security.SecurityMiddleware",
     "whitenoise.middleware.WhiteNoiseMiddleware",
     "django.contrib.sessions.middleware.SessionMiddleware",
@@ -62,7 +62,7 @@
     "core.middleware.ServiceMiddleware",
     "codecov_auth.middleware.CurrentOwnerMiddleware",
     "codecov_auth.middleware.ImpersonationMiddleware",
-    "django_prometheus.middleware.PrometheusAfterMiddleware",
+    "core.middleware.AppMetricsAfterMiddlewareWithUA",
 ]
 
 ROOT_URLCONF = "codecov.urls"

core/middleware.py

@@ -3,9 +3,34 @@
 from django.http import HttpRequest
 from django.urls import resolve
 from django.utils.deprecation import MiddlewareMixin
+from django_prometheus.middleware import (
+    Metrics,
+    PrometheusAfterMiddleware,
+    PrometheusBeforeMiddleware,
+)
 
 from utils.services import get_long_service_name
 
+# Prometheus metrics that will be annotated with User-Agent http header as label
+USER_AGENT_METRICS = [
+    "django_http_requests_unknown_latency_including_middlewares_total",
+    "django_http_requests_latency_seconds_by_view_method",
+    "django_http_requests_unknown_latency_total",
+    "django_http_ajax_requests_total",
+    "django_http_requests_total_by_method",
+    "django_http_requests_total_by_transport",
+    "django_http_requests_total_by_view_transport_method",
+    "django_http_requests_body_total_bytes",
+    "django_http_responses_total_by_templatename",
+    "django_http_responses_total_by_status",
+    "django_http_responses_total_by_status_view_method",
+    "django_http_responses_body_total_bytes",
+    "django_http_responses_total_by_charset",
+    "django_http_responses_streaming_total",
+    "django_http_exceptions_total_by_type",
+    "django_http_exceptions_total_by_view",
+]
+
 
 def get_service_long_name(request: HttpRequest) -> Optional[str]:
     resolver_match = resolve(request.path_info)
@@ -23,3 +48,39 @@ def process_view(self, request, view_func, view_args, view_kwargs):
         if service:
             view_kwargs["service"] = service
         return None
+
+
+class CustomMetricsWithUA(Metrics):
+    """
+    django_prometheus Metrics class but with extra user_agent label for applicable metrics
+    """
+
+    def register_metric(self, metric_cls, name, documentation, labelnames=(), **kwargs):
+        if name in USER_AGENT_METRICS:
+            labelnames = list(labelnames) + ["user_agent"]
+        return super().register_metric(
+            metric_cls, name, documentation, labelnames=labelnames, **kwargs
+        )
+
+
+class AppMetricsBeforeMiddlewareWithUA(PrometheusBeforeMiddleware):
+    """
+    django_prometheus monitoring middleware using custom Metrics class
+    """
+
+    metrics_cls = CustomMetricsWithUA
+
+
+class AppMetricsAfterMiddlewareWithUA(PrometheusAfterMiddleware):
+    """
+    django_prometheus monitoring middleware using custom Metrics class that injects User-Agent label when possible
+    """
+
+    metrics_cls = CustomMetricsWithUA
+
+    def label_metric(self, metric, request, response=None, **labels):
+        new_labels = labels
+        if metric._name in USER_AGENT_METRICS:
+            new_labels = {"user_agent": request.headers.get("User-Agent", "none")}
+            new_labels.update(labels)
+        return super().label_metric(metric, request, response=response, **new_labels)

core/tests/test_middleware.py

@@ -0,0 +1,131 @@
+from django.test import TestCase
+from prometheus_client import REGISTRY
+
+from core.middleware import USER_AGENT_METRICS
+
+
+# TODO: consolidate with worker/helpers/tests/unit/test_checkpoint_logger.py into shared repo
+class CounterAssertion:
+    def __init__(self, metric, labels, expected_value):
+        self.metric = metric
+        self.labels = labels
+        self.expected_value = expected_value
+
+        self.before_value = None
+        self.after_value = None
+
+    def __repr__(self):
+        return f"<CounterAssertion: {self.metric} {self.labels}>"
+
+
+# TODO: consolidate with worker/helpers/tests/unit/test_checkpoint_logger.py into shared repo
+class CounterAssertionSet:
+    def __init__(self, counter_assertions):
+        self.counter_assertions = counter_assertions
+
+    def __enter__(self):
+        for assertion in self.counter_assertions:
+            assertion.before_value = (
+                REGISTRY.get_sample_value(assertion.metric, labels=assertion.labels)
+                or 0
+            )
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        for assertion in self.counter_assertions:
+            assertion.after_value = (
+                REGISTRY.get_sample_value(assertion.metric, labels=assertion.labels)
+                or 0
+            )
+            assert (
+                assertion.after_value - assertion.before_value
+                == assertion.expected_value
+            )
+
+
+class PrometheusUserAgentLabelTest(TestCase):
+    def test_user_agent_label_added(self):
+        user_agent = "iphone"
+
+        counter_assertions = [
+            CounterAssertion(
+                "django_http_requests_latency_seconds_by_view_method_count",
+                {
+                    "view": "codecov.views.health",
+                    "method": "GET",
+                    "user_agent": user_agent,
+                },
+                1,
+            ),
+            CounterAssertion(
+                "django_http_requests_total_by_method_total",
+                {"user_agent": user_agent, "method": "GET"},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_requests_total_by_transport_total",
+                {"transport": "http", "user_agent": user_agent},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_requests_total_by_view_transport_method_total",
+                {
+                    "view": "codecov.views.health",
+                    "transport": "http",
+                    "method": "GET",
+                    "user_agent": user_agent,
+                },
+                1,
+            ),
+            CounterAssertion(
+                "django_http_requests_body_total_bytes_count",
+                {"user_agent": user_agent},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_requests_total_by_transport_total",
+                {"transport": "http", "user_agent": user_agent},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_responses_total_by_status_total",
+                {"status": "200", "user_agent": user_agent},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_responses_total_by_status_view_method_total",
+                {
+                    "status": "200",
+                    "view": "codecov.views.health",
+                    "method": "GET",
+                    "user_agent": user_agent,
+                },
+                1,
+            ),
+            CounterAssertion(
+                "django_http_responses_body_total_bytes_count",
+                {"user_agent": user_agent},
+                1,
+            ),
+            CounterAssertion(
+                "django_http_responses_total_by_charset_total",
+                {"charset": "utf-8", "user_agent": user_agent},
+                1,
+            ),
+        ]
+
+        with CounterAssertionSet(counter_assertions):
+            self.client.get(
+                "/",
+                headers={
+                    "User-Agent": user_agent,
+                },
+            )
+
+            for metric in REGISTRY.collect():
+                if metric.name in USER_AGENT_METRICS:
+                    for sample in metric.samples:
+                        assert (
+                            sample.labels["user_agent"]
+                            == "none"  # not all requests have User-Agent header defined
+                            or sample.labels["user_agent"] == user_agent
+                        )