Merge pull request #92 from backbord/master

Optionally inject Request to exemplars function.

starlette_exporter/middleware.py

@@ -1,4 +1,5 @@
 """ Middleware for exporting Prometheus metrics using Starlette """
+import inspect
 import logging
 import re
 import time
@@ -115,6 +116,21 @@ def __init__(
 
         self.labels = OrderedDict(labels) if labels is not None else None
         self.exemplars = exemplars
+        self._exemplars_req_kw = ""
+
+        if self.exemplars:
+            # if the exemplars func has an argument annotated as Request, note its name.
+            # it will be used to inject the request when the func is called
+            exemplar_sig = inspect.signature(self.exemplars)
+            for p in exemplar_sig.parameters.values():
+                if p.annotation is Request:
+                    self._exemplars_req_kw = p.name
+                    break
+            else:
+                # if there's no parameter with a Request type annotation but there is a
+                # parameter with name "request", it will be chosen for injection
+                if "request" in exemplar_sig.parameters:
+                    self._exemplars_req_kw = "request"
 
     # Default metrics
 
@@ -372,7 +388,10 @@ async def wrapped_send(message: Message) -> None:
         # note: only used for histogram observations and counters to support exemplars
         extra = {}
         if self.exemplars:
-            extra["exemplar"] = self.exemplars()
+            exemplar_kwargs = {}
+            if self._exemplars_req_kw:
+                exemplar_kwargs[self._exemplars_req_kw] = request
+            extra["exemplar"] = self.exemplars(**exemplar_kwargs)
 
         # optional response body size metric
         if (

tests/test_middleware.py

@@ -6,6 +6,7 @@
 from starlette.applications import Starlette
 from starlette.background import BackgroundTask
 from starlette.exceptions import HTTPException
+from starlette.requests import Request
 from starlette.responses import JSONResponse, Response
 from starlette.routing import Mount, Route
 from starlette.staticfiles import StaticFiles
@@ -803,3 +804,35 @@ def exemplar_fn():
             """starlette_requests_total{app_name="starlette",method="GET",path="/200",status_code="200",test="exemplar"} 1.0 # {trace_id="abc123"}"""
             in metrics
         ), metrics
+
+    @pytest.mark.parametrize("annotated", [False, True])
+    def test_exemplar_request(self, testapp, annotated: bool) -> None:
+        """test setting exemplar with request injection"""
+
+        # create a callable that returns a label/value pair to
+        # be used as an exemplar.
+        if annotated:
+
+            def exemplar_fn(r: Request):
+                return {"trace_id": r.headers.get("trace-id", "")}
+
+        else:
+
+            def exemplar_fn(request):
+                return {"trace_id": request.headers.get("trace-id", "")}
+
+        # create a label for this test so we have a unique output line
+        labels = {"test": "exemplar"}
+
+        client = TestClient(testapp(exemplars=exemplar_fn, labels=labels))
+        client.get("/200", headers={"Trace-ID": "abc123"})
+
+        metrics = client.get(
+            "/openmetrics",
+            headers={"Accept": "application/openmetrics-text", "Trace-ID": "abc123"},
+        ).content.decode()
+
+        assert (
+            """starlette_requests_total{app_name="starlette",method="GET",path="/200",status_code="200",test="exemplar"} 1.0 # {trace_id="abc123"}"""
+            in metrics
+        ), metrics