From a79946fbc1e758b877cf767782084d86d2c3b16c Mon Sep 17 00:00:00 2001
From: Roberto Polli <robipolli@gmail.com>
Date: Mon, 4 Dec 2017 18:52:36 +0100
Subject: [PATCH] Fix: #222. Support custom registry for MetricsHandler.

  Pass a registry to start_http_server via a MetricsHandler
    factory implemented as a staticmethod of the existing class.
---
 prometheus_client/exposition.py   | 23 ++++++++++++++++++++---
 prometheus_client/multiprocess.py |  1 +
 tests/test_exposition.py          |  6 +++++-
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
index 2d54e451..b79b0d07 100644
--- a/prometheus_client/exposition.py
+++ b/prometheus_client/exposition.py
@@ -83,9 +83,10 @@ def generate_latest(registry=core.REGISTRY):
 
 class MetricsHandler(BaseHTTPRequestHandler):
     """HTTP handler that gives metrics from ``core.REGISTRY``."""
+    registry = core.REGISTRY
 
     def do_GET(self):
-        registry = core.REGISTRY
+        registry = self.registry
         params = parse_qs(urlparse(self.path).query)
         if 'name[]' in params:
             registry = registry.restricted_registry(params['name[]'])
@@ -102,14 +103,30 @@ def do_GET(self):
     def log_message(self, format, *args):
         """Log nothing."""
 
+    @staticmethod
+    def factory(registry):
+        """Returns a dynamic MetricsHandler class tied
+           to the passed registry.
+        """
+        # This implementation relies on MetricsHandler.registry
+        #  (defined above and defaulted to core.REGISTRY).
+
+        # As we have unicode_literals, we need to create a str()
+        #  object for type().
+        cls_name = str('MetricsHandler')
+        MyMetricsHandler = type(cls_name, (MetricsHandler, object),
+                                {"registry": registry})
+        return MyMetricsHandler
+
 
 class _ThreadingSimpleServer(ThreadingMixIn, HTTPServer):
     """Thread per request HTTP server."""
 
 
-def start_http_server(port, addr=''):
+def start_http_server(port, addr='', registry=core.REGISTRY):
     """Starts an HTTP server for prometheus metrics as a daemon thread"""
-    httpd = _ThreadingSimpleServer((addr, port), MetricsHandler)
+    CustomMetricsHandler = MetricsHandler.factory(registry)
+    httpd = _ThreadingSimpleServer((addr, port), CustomMetricsHandler)
     t = threading.Thread(target=httpd.serve_forever)
     t.daemon = True
     t.start()
diff --git a/prometheus_client/multiprocess.py b/prometheus_client/multiprocess.py
index c097ced7..6a878b53 100644
--- a/prometheus_client/multiprocess.py
+++ b/prometheus_client/multiprocess.py
@@ -9,6 +9,7 @@
 
 from . import core
 
+
 class MultiProcessCollector(object):
     """Collector for files for multi-process mode."""
     def __init__(self, registry, path=None):
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
index 3b1fb16f..a4cc385a 100644
--- a/tests/test_exposition.py
+++ b/tests/test_exposition.py
@@ -13,7 +13,8 @@
 from prometheus_client import CollectorRegistry, generate_latest
 from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
 from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
-from prometheus_client.exposition import default_handler, basic_auth_handler
+from prometheus_client import core
+from prometheus_client.exposition import default_handler, basic_auth_handler, MetricsHandler
 
 try:
     from BaseHTTPServer import BaseHTTPRequestHandler
@@ -195,6 +196,9 @@ def my_auth_handler(url, method, timeout, headers, data):
     def test_instance_ip_grouping_key(self):
         self.assertTrue('' != instance_ip_grouping_key()['instance'])
 
+    def test_metrics_handler(self):
+        MyHandler = MetricsHandler.factory(core.REGISTRY)
+
 
 if __name__ == '__main__':
     unittest.main()