Repository: prometheus/client_python
Pull request #730 (2021-12-01 10:58:04)

Allow to add labels inside a context manager
This is an attempt to use `time()` as a context manager, but still be able to add labels with values that depend on the result of the timed operation. In our case we'd like to monitor the duration of a number of web requests, but still distinguish/label them by the returned HTTP error code. It didn't seem possible to use the provided `time()` helper in this case, but instead we would have to measure the time ourselves and then use something like `metric.labels(status=status_code).observe(duration)`.

Being able to add the label inside the context manager seemed a bit more convenient:
```python
from prometheus_client import Histogram
from requests import get

teapot = Histogram('teapot', 'A teapot', ['status'])
with teapot.time() as metric:
    response = get('https://httpbin.org/status/418')
    metric.labels(status=response.status_code)
```
This should now work with all three `time()` helpers. However, the 'observability' checks for 'gauge' and 'summary' metrics had to be deferred (to their respective `__exit__` methods; see 2b57f24c5f6879090b49443152675fd807951230), because there is no way to tell upfront if a label will in fact be added, making the metric 'observable'.

I hope that's an acceptable change. If so, the `track_inprogress()` helper should probably also be changed accordingly.

witsch: > Thanks for the PR! I think the new behavior is certainly helpful, but I am not sure of the implementation yet. Happy to brainstorm on some ways to avoid `__self__` and `__name__` if possible.

I suppose an easy change would be to pass both the metric instance and the name of the callback method when initializing `Timer`, i.e. `return Timer(self, 'observe')` instead of `return Timer(self.observe)` as it is now. This would only affect three places in the code (afaics), but I was trying to keep the changes at a minimum.

Alternatively, an alias like `callback = observe` (in `Summary` and `Histogram`) and `callback = set` (in `Gauge`) would also work. `Timer` would only have to remember the metric instance and call `.callback(…)` on it.

Then again, it's only names. Slightly funny looking perhaps, but we use `__init__` and `__call__` all the time as well… 😉
csmarchbanks: I definitely think passing the metric instance is nicer, and will work nicely with types hints as they are added. Perhaps it should be called with `return Timer(self, self.observe)`, that way it is clear we are passing in both a metric and a callback?
witsch: > I definitely think passing the metric instance is nicer, and will work nicely with types hints as they are added. Perhaps it should be called with `return Timer(self, self.observe)`, that way it is clear we are passing in both a metric and a callback?

Agreed, but I think passing `self.observe` instead of the callback name would still require us to keep [`Timer.labels()` ](https://github.com/prometheus/client_python/pull/730/files#diff-9047eea4fec76c13f77a7c5ce510635229e51349ff6ac852e30d22ab0cb88e55R64) almost the same, except for the `self._callback.__self__`, of course. We'd still have to `getattr(…, self._callback.__name__)`, though.

If that's okay with you, I can happily adjust the code tomorrow? 🙂 

csmarchbanks: Ahh, right that makes sense. I think the callback name for now sounds good, we can always change it if it becomes a problem to maintain and the new functionality will be nice.
witsch: > If that's okay with you, I can happily adjust the code tomorrow? 🙂

Yesterday was too "full", but I think now it's really a bit clearer. And I suppose I should configure my git to sign off by default… 😅
witsch: > You mentioned updating `track_inprogress` as well, I don't think that would make sense as we `inc` on `__enter__` which means we would need to know any labels upfront. Am I missing something though?

True. I had just spotted that (other) context manager, but not looked into the details yet. But yes, that wouldn't work, of course.

Patch:
diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py
index 680c916..b229b17 100644
--- a/prometheus_client/context_managers.py
+++ b/prometheus_client/context_managers.py
@@ -53,19 +53,25 @@ class InprogressTracker:
 
 
 class Timer:
-    def __init__(self, callback):
-        self._callback = callback
+    def __init__(self, metric, callback_name):
+        self._metric = metric
+        self._callback_name = callback_name
 
     def _new_timer(self):
-        return self.__class__(self._callback)
+        return self.__class__(self._metric, self._callback_name)
 
     def __enter__(self):
         self._start = default_timer()
+        return self
 
     def __exit__(self, typ, value, traceback):
         # Time can go backwards.
         duration = max(default_timer() - self._start, 0)
-        self._callback(duration)
+        callback = getattr(self._metric, self._callback_name)
+        callback(duration)
+
+    def labels(self, *args, **kw):
+        self._metric = self._metric.labels(*args, **kw)
 
     def __call__(self, f):
         def wrapped(func, *args, **kwargs):
diff --git a/prometheus_client/metrics.py b/prometheus_client/metrics.py
index 78f0adc..12f793c 100644
--- a/prometheus_client/metrics.py
+++ b/prometheus_client/metrics.py
@@ -402,8 +402,7 @@ class Gauge(MetricWrapperBase):
 
         Can be used as a function decorator or context manager.
         """
-        self._raise_if_not_observable()
-        return Timer(self.set)
+        return Timer(self, 'set')
 
     def set_function(self, f):
         """Call the provided function to return the Gauge value.
@@ -481,8 +480,7 @@ class Summary(MetricWrapperBase):
 
         Can be used as a function decorator or context manager.
         """
-        self._raise_if_not_observable()
-        return Timer(self.observe)
+        return Timer(self, 'observe')
 
     def _child_samples(self):
         return (
@@ -606,7 +604,7 @@ class Histogram(MetricWrapperBase):
 
         Can be used as a function decorator or context manager.
         """
-        return Timer(self.observe)
+        return Timer(self, 'observe')
 
     def _child_samples(self):
         samples = []


Test patch:
diff --git a/tests/test_core.py b/tests/test_core.py
index 349e475..38bc20c 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -222,13 +222,25 @@ class TestGauge(unittest.TestCase):
             time.sleep(.001)
         self.assertNotEqual(0, self.registry.get_sample_value('g'))
 
+    def test_time_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('g2', {'label1': 'foo'}))
+        with self.gauge_with_label.time() as metric:
+            metric.labels('foo')
+        self.assertLess(0, value('g2', {'label1': 'foo'}))
+
     def test_track_in_progress_not_observable(self):
         g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
         assert_not_observable(g.track_inprogress)
 
     def test_timer_not_observable(self):
         g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
-        assert_not_observable(g.time)
+
+        def manager():
+            with g.time():
+                pass
+
+        assert_not_observable(manager)
 
 
 class TestSummary(unittest.TestCase):
@@ -318,10 +330,21 @@ class TestSummary(unittest.TestCase):
             pass
         self.assertEqual(1, self.registry.get_sample_value('s_count'))
 
+    def test_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('s_with_labels_count', {'label1': 'foo'}))
+        with self.summary_with_labels.time() as metric:
+            metric.labels('foo')
+        self.assertEqual(1, value('s_with_labels_count', {'label1': 'foo'}))
+
     def test_timer_not_observable(self):
         s = Summary('test', 'help', labelnames=('label',), registry=self.registry)
 
-        assert_not_observable(s.time)
+        def manager():
+            with s.time():
+                pass
+
+        assert_not_observable(manager)
 
 
 class TestHistogram(unittest.TestCase):
@@ -435,6 +458,15 @@ class TestHistogram(unittest.TestCase):
         self.assertEqual(1, self.registry.get_sample_value('h_count'))
         self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
 
+    def test_block_decorator_with_label(self):
+        value = self.registry.get_sample_value
+        self.assertEqual(None, value('hl_count', {'l': 'a'}))
+        self.assertEqual(None, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
+        with self.labels.time() as metric:
+            metric.labels('a')
+        self.assertEqual(1, value('hl_count', {'l': 'a'}))
+        self.assertEqual(1, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
+
     def test_exemplar_invalid_label_name(self):
         self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={':o)': 'smile'})
         self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={'1': 'number'})