From 3ef865e1cccae66f63ae764762a700c5775a5190 Mon Sep 17 00:00:00 2001 From: Andreas Zeidler Date: Wed, 15 Dec 2021 20:01:30 +0100 Subject: [PATCH] Allow to add labels inside a context manager (#730) * Allow to add labels inside a context manager This way labels that depend on the result of the measured operation can be added more conveniently, e.g. the status code of an http request: 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) Signed-off-by: Andreas Zeidler * Also allow to add deferred labels for 'gauge' and 'summary' metrics For this to work the 'observability' check needs to be deferred as well, in case a label is added inside the context manager thereby making the metric observable. Signed-off-by: Andreas Zeidler * Pass metric instance and callback name to `Timer` This should make the code slightly more readable. Signed-off-by: Andreas Zeidler * Remove redundant check for observability The callbacks are already taking care of this anyway. Signed-off-by: Andreas Zeidler --- prometheus_client/context_managers.py | 14 ++++++++--- prometheus_client/metrics.py | 8 +++--- tests/test_core.py | 36 +++++++++++++++++++++++++-- 3 files changed, 47 insertions(+), 11 deletions(-) diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py index 680c9163..b229b17d 100644 --- a/prometheus_client/context_managers.py +++ b/prometheus_client/context_managers.py @@ -53,19 +53,25 @@ def wrapped(func, *args, **kwargs): 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 78f0adc7..12f793c2 100644 --- a/prometheus_client/metrics.py +++ b/prometheus_client/metrics.py @@ -402,8 +402,7 @@ def time(self): 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 @@ def time(self): 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 @@ def time(self): Can be used as a function decorator or context manager. """ - return Timer(self.observe) + return Timer(self, 'observe') def _child_samples(self): samples = [] diff --git a/tests/test_core.py b/tests/test_core.py index 349e4753..38bc20c4 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -222,13 +222,25 @@ def test_time_block_decorator(self): 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 @@ def test_block_decorator(self): 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 @@ def test_block_decorator(self): 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'})