Repository: stephenhillier/starlette_exporter (language: python)

Issue:
`from_header` but for `Response`?
Is there any way to create a label based on response in `starlette_exporter`?
My use-case is adding a hit-or-miss={hit,miss} label based on the `x-fastapi-cache` header.


Discussion / hints:
@lainiwa unfortunately, not at the moment. The custom labels are evaluated before the endpoint handler is called. Interesting idea- I'll report back if I can think of a good solution.





Fix patch:
diff --git a/README.md b/README.md
index ecc4776..1815f7b 100644
--- a/README.md
+++ b/README.md
@@ -135,13 +135,22 @@ values are constrained (see [this writeup from Grafana on cardinality](https://g
 
 ### Label helpers
 
-**`from_header(key: string, allowed_values: Optional[Iterable])`**: a convenience function for using a header value as a label.
+**`from_header(key: string, allowed_values: Optional[Iterable] = None, default: str = "")`**: a convenience function for using a header value as a label.
 
 `allowed_values` allows you to supply a list of allowed values. If supplied, header values not in the list will result in
 an empty string being returned. This allows you to constrain the label values, reducing the risk of excessive cardinality.
 
+`default`: the default value if the header does not exist.
+
 Do not use headers that could contain unconstrained values (e.g. user id) or user-supplied values.
 
+
+**`from_response_header(key: str, allowed_values: Optional[Iterable] = None, default: str = "")`**: a helper
+function that extracts a value from a response header. This may be useful if you are using a middleware
+or decorator that populates a header.
+
+The same options (and warnings) apply as the `from_header` function.
+
 ```python
 from starlette_exporter import PrometheusMiddleware, from_header
 
@@ -149,6 +158,7 @@ app.add_middleware(
   PrometheusMiddleware,
   labels={
       "host": from_header("X-Internal-Org", allowed_values=("accounting", "marketing", "product"))
+      "cache": from_response_header("X-FastAPI-Cache", allowed_values=("hit", "miss"))
     }
 ```
 
diff --git a/setup.py b/setup.py
index 3eebc55..4329e58 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="starlette_exporter",
-    version="0.21.0",
+    version="0.22.0",
     author="Stephen Hillier",
     author_email="stephenhillier@gmail.com",
     packages=["starlette_exporter"],
diff --git a/starlette_exporter/__init__.py b/starlette_exporter/__init__.py
index 2bf11e6..adb9e2d 100644
--- a/starlette_exporter/__init__.py
+++ b/starlette_exporter/__init__.py
@@ -1,6 +1,7 @@
 __all__ = [
     'PrometheusMiddleware',
     'from_header',
+    'from_response_header',
     'handle_metrics',
 ]
 
@@ -20,7 +21,7 @@
 from starlette.responses import Response
 
 from .middleware import PrometheusMiddleware
-from .labels import from_header
+from .labels import from_header, from_response_header
 
 
 def handle_metrics(request: Request) -> Response:
diff --git a/starlette_exporter/labels.py b/starlette_exporter/labels.py
index 5fcd880..1701638 100644
--- a/starlette_exporter/labels.py
+++ b/starlette_exporter/labels.py
@@ -1,7 +1,27 @@
 """utilities for working with labels"""
-from typing import Callable, Iterable, Optional
+from typing import Any, Callable, Iterable, Optional, Dict
 
 from starlette.requests import Request
+from starlette.types import Message
+
+
+class ResponseHeaderLabel:
+    """ResponseHeaderLabel is a special class that allows populating a label
+    value based on response headers. starlette_exporter will recognize that this
+    should not be called until response headers are written."""
+
+    def __init__(
+        self, key: str, allowed_values: Optional[Iterable] = None, default: str = ""
+    ) -> None:
+        self.key = key
+        self.default = default
+        self.allowed_values = allowed_values
+
+    def __call__(self, headers: Dict) -> Any:
+        v = headers.get(self.key, self.default)
+        if self.allowed_values is not None and v not in self.allowed_values:
+            return self.default
+        return v
 
 
 def from_header(key: str, allowed_values: Optional[Iterable] = None) -> Callable:
@@ -36,3 +56,12 @@ def inner(r: Request):
         return v
 
     return inner
+
+
+def from_response_header(
+    key: str, allowed_values: Optional[Iterable] = None, default: str = ""
+):
+    """returns a callable class that retrieves a header value from response headers.
+    starlette_exporter will automatically populate this label value when response headers
+    are written."""
+    return ResponseHeaderLabel(key, allowed_values, default)
diff --git a/starlette_exporter/middleware.py b/starlette_exporter/middleware.py
index cca8c4f..e71d4b7 100644
--- a/starlette_exporter/middleware.py
+++ b/starlette_exporter/middleware.py
@@ -24,6 +24,8 @@
 from starlette.routing import BaseRoute, Match
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
+from starlette_exporter.labels import ResponseHeaderLabel
+
 from . import optional_metrics
 
 logger = logging.getLogger("starlette_exporter")
@@ -114,7 +116,6 @@ def __init__(
             self.optional_metrics_list = optional_metrics
         self.always_use_int_status = always_use_int_status
 
-        self.labels = OrderedDict(labels) if labels is not None else None
         self.exemplars = exemplars
         self._exemplars_req_kw = ""
 
@@ -132,7 +133,22 @@ def __init__(
                 if "request" in exemplar_sig.parameters:
                     self._exemplars_req_kw = "request"
 
+
+        # split labels into request and response labels.
+        # response labels will be evaluated while the response is
+        # written.
+        self.request_labels = OrderedDict({})
+        self.response_labels: OrderedDict[str, ResponseHeaderLabel] = OrderedDict({})
+
+        if labels is not None:
+            for k, v in labels.items():
+                if isinstance(v, ResponseHeaderLabel):
+                    self.response_labels[k] = v
+                else:
+                    self.request_labels[k] = v
+
     # Default metrics
+    # Starlette initialises middleware multiple times, so store metrics on the class
 
     @property
     def request_count(self):
@@ -146,7 +162,8 @@ def request_count(self):
                     "path",
                     "status_code",
                     "app_name",
-                    *self._default_label_keys(),
+                    *self.request_labels.keys(),
+                    *self.response_labels.keys(),
                 ),
             )
         return PrometheusMiddleware._metrics[metric_name]
@@ -174,7 +191,8 @@ def response_body_size_count(self):
                         "path",
                         "status_code",
                         "app_name",
-                        *self._default_label_keys(),
+                        *self.request_labels.keys(),
+                        *self.response_labels.keys(),
                     ),
                 )
             return PrometheusMiddleware._metrics[metric_name]
@@ -200,7 +218,8 @@ def request_body_size_count(self):
                         "path",
                         "status_code",
                         "app_name",
-                        *self._default_label_keys(),
+                        *self.request_labels.keys(),
+                        *self.response_labels.keys(),
                     ),
                 )
             return PrometheusMiddleware._metrics[metric_name]
@@ -219,7 +238,8 @@ def request_time(self):
                     "path",
                     "status_code",
                     "app_name",
-                    *self._default_label_keys(),
+                    *self.request_labels.keys(),
+                    *self.response_labels.keys(),
                 ),
                 **self.kwargs,
             )
@@ -232,23 +252,15 @@ def requests_in_progress(self):
             PrometheusMiddleware._metrics[metric_name] = Gauge(
                 metric_name,
                 "Total HTTP requests currently in progress",
-                ("method", "app_name", *self._default_label_keys()),
+                ("method", "app_name", *self.request_labels.keys()),
                 multiprocess_mode="livesum",
             )
         return PrometheusMiddleware._metrics[metric_name]
 
-    def _default_label_keys(self) -> List[str]:
-        if self.labels is None:
-            return []
-        return list(self.labels.keys())
-
-    async def _default_label_values(self, request: Request):
-        if self.labels is None:
-            return []
-
+    async def _request_label_values(self, request: Request) -> List[str]:
         values: List[str] = []
 
-        for k, v in self.labels.items():
+        for k, v in self.request_labels.items():
             if callable(v):
                 parsed_value = ""
                 # if provided a callable, try to use it on the request.
@@ -267,6 +279,34 @@ async def _default_label_values(self, request: Request):
 
         return values
 
+    def _response_label_values(self, message: Message) -> List[str]:
+        values: List[str] = []
+
+        # bail if no response labels were defined by the user
+        if not self.response_labels:
+            return values
+
+        # create a dict of headers to make it easy to find keys
+        headers = {
+            k.decode("utf-8"): v.decode("utf-8")
+            for (k, v) in message.get("headers", ())
+        }
+
+        for k, v in self.response_labels.items():
+            # currently only ResponseHeaderLabel supported
+            if isinstance(v, ResponseHeaderLabel):
+                parsed_value = ""
+                try:
+                    result = v(headers)
+                except Exception:
+                    logger.warn(f"label function for {k} failed", exc_info=True)
+                else:
+                    parsed_value = str(result)
+                values.append(parsed_value)
+
+
+        return values
+
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if scope["type"] not in ["http"]:
             await self.app(scope, receive, send)
@@ -288,13 +328,16 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         begin = time.perf_counter()
         end = None
 
-        default_labels = await self._default_label_values(request)
+        request_labels = await self._request_label_values(request)
 
         # Increment requests_in_progress gauge when request comes in
-        self.requests_in_progress.labels(method, self.app_name, *default_labels).inc()
+        self.requests_in_progress.labels(method, self.app_name, *request_labels).inc()
 
         status_code = None
 
+        # custom response label values, to be populated when response is written.
+        response_labels = []
+
         # optional request and response body size metrics
         response_body_size: int = 0
 
@@ -311,10 +354,13 @@ async def wrapped_send(message: Message) -> None:
                 nonlocal status_code
                 status_code = message["status"]
 
+                nonlocal response_labels
+                response_labels = self._response_label_values(message)
+
                 if self.always_use_int_status:
                     try:
                         status_code = int(message["status"])
-                    except ValueError as e:
+                    except ValueError:
                         logger.warning(
                             f"always_use_int_status flag selected but failed to convert status_code to int for value: {status_code}"
                         )
@@ -348,11 +394,18 @@ async def wrapped_send(message: Message) -> None:
         except Exception as e:
             status_code = 500
             exception = e
+        finally:
+            # Decrement 'requests_in_progress' gauge after response sent
+            self.requests_in_progress.labels(
+                method, self.app_name, *request_labels
+            ).dec()
 
-        # Decrement 'requests_in_progress' gauge after response sent
-        self.requests_in_progress.labels(
-            method, self.app_name, *default_labels
-        ).dec()
+        if status_code is None:
+            if await request.is_disconnected():
+                # In case no response was returned and the client is disconnected, 499 is reported as status code.
+                status_code = 499
+            else:
+                status_code = 500
 
         if self.filter_unhandled_paths or self.group_paths:
             grouped_path: Optional[str] = None
@@ -369,20 +422,21 @@ async def wrapped_send(message: Message) -> None:
                     raise exception
                 return
 
+
             # group_paths enables returning the original router path (with url param names)
             # for example, when using this option, requests to /api/product/1 and /api/product/3
             # will both be grouped under /api/product/{product_id}. See the README for more info.
             if self.group_paths and grouped_path is not None:
                 path = grouped_path
 
-        if status_code is None:
-            if await request.is_disconnected():
-                # In case no response was returned and the client is disconnected, 499 is reported as status code.
-                status_code = 499
-            else:
-                status_code = 500
-
-        labels = [method, path, status_code, self.app_name, *default_labels]
+        labels = [
+                method,
+                path,
+                status_code,
+                self.app_name,
+                *request_labels,
+                *response_labels,
+            ]
 
         # optional extra arguments to be passed as kwargs to observations
         # note: only used for histogram observations and counters to support exemplars


Test patch:
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
index 21381b5..eee6568 100644
--- a/tests/test_middleware.py
+++ b/tests/test_middleware.py
@@ -16,6 +16,7 @@
 from starlette_exporter import (
     PrometheusMiddleware,
     from_header,
+    from_response_header,
     handle_metrics,
     handle_openmetrics,
 )
@@ -41,8 +42,8 @@ def _testapp(**middleware_options):
         app.add_route("/metrics", handle_metrics)
         app.add_route("/openmetrics", handle_openmetrics)
 
-        def normal_response(request):
-            return JSONResponse({"message": "Hello World"})
+        def normal_response(_):
+            return JSONResponse({"message": "Hello World"}, headers={"foo": "baz"})
 
         app.add_route("/200", normal_response, methods=["GET", "POST", "OPTIONS"])
         app.add_route(
@@ -302,7 +303,7 @@ def test_mounted_path_404(self, testapp):
 
     def test_mounted_path_404_filter(self, testapp):
         """test an unhandled path from mounted base path can be excluded from metrics"""
-        client = TestClient(testapp(filter_unhandled_paths=True))
+        client = TestClient(testapp())
         client.get("/mounted/404")
         metrics = client.get("/metrics").content.decode()
 
@@ -778,6 +779,72 @@ def test_from_header_allowed_values_disallowed_value(self, testapp):
             in metrics
         ), metrics
 
+    def test_from_response_header(self, testapp):
+        """test with the library-provided from_response_header function"""
+        labels = {"foo": from_response_header("foo"), "hello": "world"}
+        client = TestClient(testapp(labels=labels))
+        client.get("/200")
+        metrics = client.get("/metrics").content.decode()
+
+        assert (
+            """starlette_requests_total{app_name="starlette",foo="baz",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
+            in metrics
+        ), metrics
+
+    def test_from_response_header_default(self, testapp):
+        """test with the library-provided from_response_header function, with a missing header
+        (testing the default value)"""
+        labels = {
+            "foo": from_response_header("x-missing-header", default="yo"),
+            "hello": "world",
+        }
+        client = TestClient(testapp(labels=labels))
+        client.get("/200")
+        metrics = client.get("/metrics").content.decode()
+
+        assert (
+            """starlette_requests_total{app_name="starlette",foo="yo",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
+            in metrics
+        ), metrics
+
+    def test_from_response_header_allowed_values(self, testapp):
+        """test with the library-provided from_response_header function, with a header value
+        that matches an allowed value."""
+        labels = {
+            "foo": from_response_header("foo", allowed_values=("bar", "baz")),
+            "hello": "world",
+        }
+        client = TestClient(testapp(labels=labels))
+        client.get("/200")
+        metrics = client.get("/metrics").content.decode()
+
+        assert (
+            """starlette_requests_total{app_name="starlette",foo="baz",hello="world",method="GET",path="/200",status_code="200"} 1.0"""
+            in metrics
+        ), metrics
+
+    def test_from_response_header_allowed_values_disallowed(self, testapp):
+        """test with the library-provided from_response_header function, with a header
+        value that does not match any of the allowed_values"""
+
+        labels = {
+            "foo": from_response_header("foo", allowed_values=("bar", "bam")),
+            "hello": "world",
+        }
+        client = TestClient(testapp(labels=labels))
+        client.get("/200", headers={"foo": "zounds"})
+        metrics = client.get("/metrics").content.decode()
+
+        assert (
+            """starlette_requests_total{app_name="starlette",foo="baz",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"""
+            in metrics
+        ), metrics
+
 
 class TestExemplars:
     """tests for adding an exemplar to the histogram and counters"""