From 026301ad4bb1e6c5a1659f7e1ad58ebc6f1d086e Mon Sep 17 00:00:00 2001 From: Will Newton Date: Wed, 22 Dec 2021 11:03:33 +0000 Subject: [PATCH 1/2] Allow adding arbitrary labels to metrics when generating output It's useful to be able to apply a label across all metrics, for example when using a multi-processing framework to apply a label for worker ID to all metrics (including from builtin collector types). This change adds a new Registry type that applies a dictionary of labels when metrics are collected. Signed-off-by: Will Newton --- prometheus_client/registry.py | 25 +++++++++++++++++++++++++ tests/test_core.py | 7 +++++++ 2 files changed, 32 insertions(+) diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py index fe435cd1..38d5f478 100644 --- a/prometheus_client/registry.py +++ b/prometheus_client/registry.py @@ -82,6 +82,18 @@ def collect(self): for collector in collectors: yield from collector.collect() + def labelled_registry(self, extra_labels): + """Returns object that collects metrics with additional labels. + + Returns an object which upon collect() will return + samples with additional labels given in the dictionary. + + Intended usage is: + generate_latest(REGISTRY.labelled_registry(['a_timeseries'])) + + Experimental.""" + return LabelledRegistry(extra_labels, self) + def restricted_registry(self, names): """Returns object that only collects some metrics. @@ -128,6 +140,19 @@ def get_sample_value(self, name, labels=None): return None +class LabelledRegistry: + """Metric collector registry that applies extra labels to metrics.""" + + def __init__(self, extra_labels, registry): + self._extra_labels = extra_labels + self._registry = registry + + def collect(self): + for metric in self._registry.collect(): + metric.samples = [s._replace(labels = {**s.labels, **self._extra_labels}) for s in metric.samples] + yield metric + + class RestrictedRegistry: def __init__(self, names, registry): self._name_set = set(names) diff --git a/tests/test_core.py b/tests/test_core.py index 38bc20c4..54d6a3d7 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -929,6 +929,13 @@ def test_restricted_registry_does_not_yield_while_locked(self): for _ in registry.restricted_registry(['target_info', 's_sum']).collect(): self.assertFalse(registry._lock.locked()) + def test_labelled_registry(self): + extra_labels = {'label1': 'value1'} + registry = CollectorRegistry(extra_labels) + Counter('c_total', 'help', registry=registry) + metrics = list(registry.labelled_registry(extra_labels).collect()) + self.assertEqual(metrics[0].samples[0].labels, extra_labels) + if __name__ == '__main__': unittest.main() From 3ec8e22b9e46c96af37155fd016dfe9318fb5174 Mon Sep 17 00:00:00 2001 From: Will Newton Date: Wed, 22 Dec 2021 11:22:22 +0000 Subject: [PATCH 2/2] prometheus_client/registry.py: Fix flake8 errors Signed-off-by: Will Newton --- prometheus_client/registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py index 38d5f478..75ec7544 100644 --- a/prometheus_client/registry.py +++ b/prometheus_client/registry.py @@ -149,7 +149,7 @@ def __init__(self, extra_labels, registry): def collect(self): for metric in self._registry.collect(): - metric.samples = [s._replace(labels = {**s.labels, **self._extra_labels}) for s in metric.samples] + metric.samples = [s._replace(labels={**s.labels, **self._extra_labels}) for s in metric.samples] yield metric