8000 Allow to add labels inside a context manager (#730) · christopherime/client_python@3ef865e · GitHub
[go: up one dir, main page]
Skip to content

Commit 3ef865e

Browse files
witschwitsch
authored
Allow to add labels inside a context manager (prometheus#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 <andreas.zeidler@zeit.de> * 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 <andreas.zeidler@zeit.de> * Pass metric instance and callback name to `Timer` This should make the code slightly more readable. Signed-off-by: Andreas Zeidler <andreas.zeidler@zeit.de> * Remove redundant check for observability The callbacks are already taking care of this anyway. Signed-off-by: Andreas Zeidler <andreas.zeidler@zeit.de>
1 parent 7c44be2 commit 3ef865e

File tree

3 files changed

+47
-11
lines changed

3 files changed

+47
-11
lines changed

‎prometheus_client/context_managers.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -53,19 +53,25 @@ def wrapped(func, *args, **kwargs):
5353

5454

5555
class Timer:
56-
def __init__(self, callback):
57-
self._callback = callback
56+
def __init__(self, metric, callback_name):
57+
self._metric = metric
58+
self._callback_name = callback_name
5859

5960
def _new_timer(self):
60-
return self.__class__(self._callback)
61+
return self.__class__(self._metric, self._callback_name)
6162

6263
def __enter__(self):
6364
self._start = default_timer()
65+
return self
6466

6567
def __exit__(self, typ, value, traceback):
6668
# Time can go backwards.
6769
duration = max(default_timer() - self._start, 0)
68-
self._callback(duration)
70+
callback = getattr(self._metric, self._callback_name)
71+
callback(duration)
72+
73+
def labels(self, *args, **kw):
74+
self._metric = self._metric.labels(*args, **kw)
6975

7076
def __call__(self, f):
7177
def wrapped(func, *args, **kwargs):

‎prometheus_client/metrics.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -402,8 +402,7 @@ def time(self):
402402
403403
Can be used as a function decorator or context manager.
404404
"""
405-
self._raise_if_not_observable()
406-
return Timer(self.set)
405+
return Timer(self, 'set')
407406

408407
def set_function(self, f):
409408
"""Call the provided function to return the Gauge value.
@@ -481,8 +480,7 @@ def time(self):
481480
482481
Can be used as a function decorator or context manager.
483482
"""
484-
self._raise_if_not_observable()
485-
return Timer(self.observe)
483+
return Timer(self, 'observe')
486484

487485
def _child_samples(self):
488486
return (
@@ -606,7 +604,7 @@ def time(self):
606604
607605
Can be used as a function decorator or context manager.
608606
"""
609-
return Timer(self.observe)
607+
return Timer(self, 'observe')
610608

611609
def _child_samples(self):
612610
samples = []

‎tests/test_core.py

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -222,13 +222,25 @@ def test_time_block_decorator(self):
222222
time.sleep(.001)
223223
self.assertNotEqual(0, self.registry.get_sample_value('g'))
224224

225+
def test_time_block_decorator_with_label(self):
226+
value = self.registry.get_sample_value
227+
self.assertEqual(None, value('g2', {'label1': 'foo'}))
228+
with self.gauge_with_label.time() as metric:
229+
metric.labels('foo')
230+
self.assertLess(0, value('g2', {'label1': 'foo'}))
231+
225232
def test_track_in_progress_not_observable(self):
226233
g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
227234
assert_not_observable(g.track_inprogress)
228235

229236
def test_timer_not_observable(self):
230237
g = Gauge('test', 'help', labelnames=('label',), registry=self.registry)
231-
assert_not_observable(g.time)
238+
239+
def manager():
240+
with g.time():
241+
pass
242+
243+
assert_not_observable(manager)
232244

233245

234246
class TestSummary(unittest.TestCase):
@@ -318,10 +330,21 @@ def test_block_decorator(self):
318330
pass
319331
self.assertEqual(1, self.registry.get_sample_value('s_count'))
320332

333+
def test_block_decorator_with_label(self):
334+
value = self.registry.get_sample_value
335+
self.assertEqual(None, value('s_with_labels_count', {'label1': 'foo'}))
336+
with self.summary_with_labels.time() as metric:
337+
metric.labels('foo')
338+
self.assertEqual(1, value('s_with_labels_count', {'label1': 'foo'}))
339+
321340
def test_timer_not_observable(self):
322341
s = Summary('test', 'help', labelnames=('label',), registry=self.registry)
323342

324-
assert_not_observable(s.time)
343+
def manager():
344+
with s.time():
345+
pass
346+
347+
assert_not_observable(manager)
325348

326349

327350
class TestHistogram(unittest.TestCase):
@@ -435,6 +458,15 @@ def test_block_decorator(self):
435458
self.assertEqual(1, self.registry.get_sample_value('h_count'))
436459
self.assertEqual(1, self.registry.get_sample_value('h_bucket', {'le': '+Inf'}))
437460

461+
def test_block_decorator_with_label(self):
462+
value = self.registry.get_sample_value
463+
self.assertEqual(None, value('hl_count', {'l': 'a'}))
464+
self.assertEqual(None, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
465+
with self.labels.time() as metric:
466+
metric.labels('a')
467+
self.assertEqual(1, value('hl_count', {'l': 'a'}))
468+
self.assertEqual(1, value('hl_bucket', {'le': '+Inf', 'l': 'a'}))
469+
438470
def test_exemplar_invalid_label_name(self):
439471
self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={':o)': 'smile'})
440472
self.assertRaises(ValueError, self.histogram.observe, 3.0, exemplar={'1': 'number'})

0 commit comments

Comments
 (0)
0