From 97f05b50a4bca229274f496e5df1508dad69e3e6 Mon Sep 17 00:00:00 2001
From: Will Newton <will.newton@gmail.com>
Date: Mon, 20 Dec 2021 09:49:44 +0000
Subject: [PATCH] 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 alows passing a dictionary of
labels and values to generate_latest which will then be applied
to all generated metrics.

Signed-off-by: Will Newton <will.newton@gmail.com>
---
 prometheus_client/exposition.py |  7 ++++---
 tests/test_exposition.py        | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
index edb7ae70..2167b63b 100644
--- a/prometheus_client/exposition.py
+++ b/prometheus_client/exposition.py
@@ -135,15 +135,16 @@ def start_wsgi_server(port, addr='', registry=REGISTRY):
 start_http_server = start_wsgi_server
 
 
-def generate_latest(registry=REGISTRY):
+def generate_latest(registry=REGISTRY, extra_labels={}):
     """Returns the metrics from the registry in latest text format as a string."""
 
     def sample_line(line):
-        if line.labels:
+        labels = {**line.labels, **extra_labels}
+        if labels:
             labelstr = '{{{0}}}'.format(','.join(
                 ['{}="{}"'.format(
                     k, v.replace('\\', r'\\').replace('\n', r'\n').replace('"', r'\"'))
-                    for k, v in sorted(line.labels.items())]))
+                    for k, v in sorted(labels.items())]))
         else:
             labelstr = ''
         timestamp = ''
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
index a52805ea..fc33bb52 100644
--- a/tests/test_exposition.py
+++ b/tests/test_exposition.py
@@ -197,6 +197,26 @@ def collect(self):
 ts{foo="f"} 0.0 123000
 """, generate_latest(self.registry))
 
+    def test_extra_labels(self):
+        c1 = Counter('c1', 'A counter', ["label1"], registry=self.registry)
+        c1.labels(label1="value1").inc()
+        c2 = Counter('c2', 'Another counter', registry=self.registry)
+        c2.inc()
+        extra_labels = {"label2": "value2"}
+        self.assertEqual(b"""# HELP c1_total A counter
+# TYPE c1_total counter
+c1_total{label1="value1",label2="value2"} 1.0
+# HELP c1_created A counter
+# TYPE c1_created gauge
+c1_created{label1="value1",label2="value2"} 123.456
+# HELP c2_total Another counter
+# TYPE c2_total counter
+c2_total{label2="value2"} 1.0
+# HELP c2_created Another counter
+# TYPE c2_created gauge
+c2_created{label2="value2"} 123.456
+""", generate_latest(self.registry, extra_labels))
+
 
 class TestPushGateway(unittest.TestCase):
     def setUp(self):