From 7c1cefd211139bcfb33cbbf59b27f46d179ec2d2 Mon Sep 17 00:00:00 2001
From: Brian Brazil <brian.brazil@robustperception.io>
Date: Tue, 13 Aug 2019 13:57:24 +0100
Subject: [PATCH] Add target_info to registries

This allows target labels as needed by push-based systems to be
provided, in a way that doesn't mess up Prometheus's own
top-down pull based approach to SD.

Signed-off-by: Brian Brazil <brian.brazil@robustperception.io>
---
 prometheus_client/registry.py | 33 +++++++++++++++++++++++++++++++--
 tests/test_core.py            | 31 +++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 2 deletions(-)

diff --git a/prometheus_client/registry.py b/prometheus_client/registry.py
index fa2717fb..7f86a000 100644
--- a/prometheus_client/registry.py
+++ b/prometheus_client/registry.py
@@ -12,11 +12,13 @@ class CollectorRegistry(object):
     exposition formats.
     """
 
-    def __init__(self, auto_describe=False):
+    def __init__(self, auto_describe=False, target_info=None):
         self._collector_to_names = {}
         self._names_to_collectors = {}
         self._auto_describe = auto_describe
         self._lock = Lock()
+        self._target_info = {}
+        self.set_target_info(target_info)
 
     def register(self, collector):
         """Add a collector to the registry."""
@@ -69,8 +71,13 @@ def _get_names(self, collector):
     def collect(self):
         """Yields metrics from the collectors in the registry."""
         collectors = None
+        ti = None
         with self._lock:
             collectors = copy.copy(self._collector_to_names)
+            if self._target_info:
+                ti = self._target_info_metric()
+        if ti:
+            yield ti
         for collector in collectors:
             for metric in collector.collect():
                 yield metric
@@ -87,11 +94,14 @@ def restricted_registry(self, names):
         Experimental."""
         names = set(names)
         collectors = set()
+        metrics = []
         with self._lock:
+            if 'target_info' in names and self._target_info:
+                metrics.append(self._target_info_metric())
+                names.remove('target_info')
             for name in names:
                 if name in self._names_to_collectors:
                     collectors.add(self._names_to_collectors[name])
-        metrics = []
         for collector in collectors:
             for metric in collector.collect():
                 samples = [s for s in metric.samples if s[0] in names]
@@ -106,6 +116,25 @@ def collect(self):
 
         return RestrictedRegistry()
 
+    def set_target_info(self, labels):
+        with self._lock:
+            if labels:
+                if not self._target_info and 'target_info' in self._names_to_collectors:
+                    raise ValueError('CollectorRegistry already contains a target_info metric')
+                self._names_to_collectors['target_info'] = None
+            elif self._target_info:
+                self._names_to_collectors.pop('target_info', None)
+            self._target_info = labels
+
+    def get_target_info(self):
+        with self._lock:
+            return self._target_info
+
+    def _target_info_metric(self):
+        m = Metric('target', 'Target metadata', 'info')
+        m.add_sample('target_info', self._target_info, 1)
+        return m
+
     def get_sample_value(self, name, labels=None):
         """Returns the sample value, or None if not found.
 
diff --git a/tests/test_core.py b/tests/test_core.py
index 9ad947cd..dac5c044 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -664,6 +664,9 @@ def test_duplicate_metrics_raises(self):
         # The name of the histogram itself isn't taken.
         Gauge('h', 'help', registry=registry)
 
+        Info('i', 'help', registry=registry)
+        self.assertRaises(ValueError, Gauge, 'i_info', 'help', registry=registry)
+
     def test_unregister_works(self):
         registry = CollectorRegistry()
         s = Summary('s', 'help', registry=registry)
@@ -696,6 +699,34 @@ def test_restricted_registry(self):
         m.samples = [Sample('s_sum', {}, 7)]
         self.assertEquals([m], registry.restricted_registry(['s_sum']).collect())
 
+    def test_target_info_injected(self):
+        registry = CollectorRegistry(target_info={'foo': 'bar'})
+        self.assertEqual(1, registry.get_sample_value('target_info', {'foo': 'bar'}))
+
+    def test_target_info_duplicate_detected(self):
+        registry = CollectorRegistry(target_info={'foo': 'bar'})
+        self.assertRaises(ValueError, Info, 'target', 'help', registry=registry)
+
+        registry.set_target_info({})
+        i = Info('target', 'help', registry=registry)
+        registry.set_target_info({})
+        self.assertRaises(ValueError, Info, 'target', 'help', registry=registry)
+        self.assertRaises(ValueError, registry.set_target_info, {'foo': 'bar'})
+        registry.unregister(i)
+        registry.set_target_info({'foo': 'bar'})
+
+    def test_target_info_restricted_registry(self):
+        registry = CollectorRegistry(target_info={'foo': 'bar'})
+        Summary('s', 'help', registry=registry).observe(7)
+
+        m = Metric('s', 'help', 'summary')
+        m.samples = [Sample('s_sum', {}, 7)]
+        self.assertEquals([m], registry.restricted_registry(['s_sum']).collect())
+
+        m = Metric('target', 'Target metadata', 'info')
+        m.samples = [Sample('target_info', {'foo': 'bar'}, 1)]
+        self.assertEquals([m], registry.restricted_registry(['target_info']).collect())
+
 
 if __name__ == '__main__':
     unittest.main()