Repository: stephenhillier/starlette_exporter

Issue:
Feature request: Support async label callback
Would be nice if we could define async callback functions that will produce metric label. If request body is needed to be parsed in the callback, we need to await since starlette body() is a coroutine. I inspected the code and it seems trivial to add. Thanks!

Discussion:

Patch:
diff --git a/starlette_exporter/middleware.py b/starlette_exporter/middleware.py
index c5ca9af..058f0e4 100644
--- a/starlette_exporter/middleware.py
+++ b/starlette_exporter/middleware.py
@@ -2,6 +2,7 @@
 from collections import OrderedDict
 import time
 import logging
+from inspect import iscoroutine
 from typing import Any, Callable, List, Mapping, Optional, ClassVar, Dict, Union, Sequence
 
 from prometheus_client import Counter, Histogram, Gauge
@@ -208,7 +209,7 @@ class PrometheusMiddleware:
             return []
         return list(self.labels.keys())
 
-    def _default_label_values(self, request: Request):
+    async def _default_label_values(self, request: Request):
         if self.labels is None:
             return []
 
@@ -220,6 +221,8 @@ class PrometheusMiddleware:
                 # if provided a callable, try to use it on the request.
                 try:
                     result = v(request)
+                    if iscoroutine(result):
+                        result = await result
                 except Exception:
                     logger.warn(f"label function for {k} failed", exc_info=True)
                 else:
@@ -248,7 +251,7 @@ class PrometheusMiddleware:
         begin = time.perf_counter()
         end = None
 
-        default_labels = self._default_label_values(request)
+        default_labels = await self._default_label_values(request)
 
         # Increment requests_in_progress gauge when request comes in
         self.requests_in_progress.labels(method, self.app_name, *default_labels).inc()


Test patch:
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
index 22059d5..1faa917 100644
--- a/tests/test_middleware.py
+++ b/tests/test_middleware.py
@@ -603,8 +603,13 @@ class TestDefaultLabels:
 
     def test_from_header_allowed_values_disallowed_value(self, testapp):
         """test with the library-provided from_header function"""
+
+        async def async_bar(request):
+            return "bar"
+
         labels = {
             "foo": from_header("foo", allowed_values=("bar", "baz")),
+            "bar": async_bar,
             "hello": "world",
         }
         client = TestClient(testapp(labels=labels))
@@ -612,12 +617,12 @@ class TestDefaultLabels:
         metrics = client.get("/metrics").content.decode()
 
         assert (
-            """starlette_requests_total{app_name="starlette",foo="zounds",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
+            """starlette_requests_total{app_name="starlette",bar="bar",foo="zounds",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
             not in metrics
         ), metrics
 
         assert (
-            """starlette_requests_total{app_name="starlette",foo="",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
+            """starlette_requests_total{app_name="starlette",bar="bar",foo="",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
             in metrics
         ), metrics