Add ability to limit what HTTP exposition returns.

brian-brazil · brian-brazil · commit 29e83adf551c · 2016-11-03T17:01:33.000Z
This works by specifiny URL parameters called name[].
Only the required collectors will be called, and
then filtered down.

This depends on the collectors sufficiently implementing
describe.
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -97,6 +97,35 @@ def collect(self):
             for metric in collector.collect():
                 yield metric
 
+    def restricted_registry(self, names):
+        '''Returns object that only collects some metrics.
+
+        Returns an object which upon collect() will return
+        only samples with the given names.
+
+        Intended usage is:
+            generate_latest(REGISTRY.restricted_registry(['a_timeseries']))
+
+        Experimental.'''
+        names = set(names)
+        collectors = set()
+        with self._lock:
+            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]
+                if samples:
+                    m = Metric(metric.name, metric.documentation, metric.type)
+                    m.samples = samples
+                    metrics.append(m)
+        class RestrictedRegistry(object):
+            def collect(self):
+                return metrics
+        return RestrictedRegistry()
+
     def get_sample_value(self, name, labels=None):
         '''Returns the sample value, or None if not found.
 
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -15,13 +15,14 @@
     from BaseHTTPServer import HTTPServer
     from urllib2 import build_opener, Request, HTTPHandler
     from urllib import quote_plus
+    from urlparse import parse_qs, urlparse
 except ImportError:
     # Python 3
     unicode = str
     from http.server import BaseHTTPRequestHandler
     from http.server import HTTPServer
     from urllib.request import build_opener, Request, HTTPHandler
-    from urllib.parse import quote_plus
+    from urllib.parse import quote_plus, parse_qs, urlparse
 
 
 CONTENT_TYPE_LATEST = str('text/plain; version=0.0.4; charset=utf-8')
@@ -34,7 +35,11 @@ def prometheus_app(environ, start_response):
         status = str('200 OK')
         headers = [(str('Content-type'), CONTENT_TYPE_LATEST)]
         start_response(status, headers)
-        return [generate_latest(registry)]
+        params = parse_qs(environ['QUERY_STRING'])
+        r = registry
+        if 'name[]' in params:
+            r = r.restricted_registry(params['name[]'])
+        return [generate_latest(r)]
     return prometheus_app
 
 
@@ -73,7 +78,11 @@ def do_GET(self):
         self.send_response(200)
         self.send_header('Content-Type', CONTENT_TYPE_LATEST)
         self.end_headers()
-        self.wfile.write(generate_latest(core.REGISTRY))
+        registry = core.REGISTRY
+        params = parse_qs(urlparse(self.path).query)
+        if 'name[]' in params:
+            registry = registry.restricted_registry(params['name[]'])
+        self.wfile.write(generate_latest(registry))
 
     def log_message(self, format, *args):
         return
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -431,6 +431,15 @@ def test_autodescribe_disabled_by_default(self):
         self.custom_collector(CounterMetricFamily('c', 'help', value=1), registry)
         self.assertRaises(ValueError, self.custom_collector, CounterMetricFamily('c', 'help', value=1), registry)
 
+    def test_restricted_registry(self):
+        registry = CollectorRegistry()
+        Counter('c', 'help', registry=registry)
+        Summary('s', 'help', registry=registry).observe(7)
+
+        m = Metric('s', 'help', 'summary')
+        m.samples = [('s_sum', {}, 7)]
+        self.assertEquals([m], registry.restricted_registry(['s_sum']).collect())
+
 
 if __name__ == '__main__':
     unittest.main()
