Use closure to pass args to handlers (refs prometheus#60).

rossigee · rossigee · commit 44df258b203e · 2017-01-25T11:58:12.000+07:00
diff --git a/README.md b/README.md
@@ -321,6 +321,24 @@ for more information.
 `instance_ip_grouping_key` returns a grouping key with the instance label set
 to the host's IP address.
 
+### Handlers for authentication
+
+If the push gateway you are connecting to is protected with HTTP Basic Auth,
+you can use a special handler to set the Authorization header.
+
+```python
+from prometheus_client import CollectorRegistry, Gauge, push_to_gateway
+import os
+
+def my_auth_handler(url, method, timeout, headers, data):
+    username = os.environ['PUSHGW_USERNAME']
+    password = os.environ['PUSHGW_PASSWORD']
+    return basic_auth_handler(url, method, timeout, headers, data, username, password)
+registry = CollectorRegistry()
+g = Gauge('job_last_success_unixtime', 'Last time a batch job successfully finished', registry=registry)
+g.set_to_current_time()
+push_to_gateway('localhost:9091', job='batchA', registry=registry, handler=my_auth_handler)
+```
 
 ## Bridges
 
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -119,7 +119,7 @@ def write_to_textfile(path, registry):
     os.rename(tmppath, path)
 
 
-def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=None, handler_args=None):
+def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=None):
     '''Push metrics to the given pushgateway.
 
     `gateway` the url for your push gateway. Either of the form
@@ -155,14 +155,13 @@ def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, han
               failure.
               'content' is the data which should be used to form the HTTP
               Message Body.
-    `handler_args` is an optional dict of extra arguments to provide
 
     This overwrites all metrics with the same job and grouping_key.
     This uses the PUT HTTP method.'''
-    _use_gateway('PUT', gateway, job, registry, grouping_key, timeout, handler, handler_args)
+    _use_gateway('PUT', gateway, job, registry, grouping_key, timeout, handler)
 
 
-def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=None, handler_args=None):
+def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=None):
     '''PushAdd metrics to the given pushgateway.
 
     `gateway` the url for your push gateway. Either of the form
@@ -183,10 +182,10 @@ def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None,
 
     This replaces metrics with the same name, job and grouping_key.
     This uses the POST HTTP method.'''
-    _use_gateway('POST', gateway, job, registry, grouping_key, timeout, handler, handler_args)
+    _use_gateway('POST', gateway, job, registry, grouping_key, timeout, handler)
 
 
-def delete_from_gateway(gateway, job, grouping_key=None, timeout=None, handler=None, handler_args=None):
+def delete_from_gateway(gateway, job, grouping_key=None, timeout=None, handler=None):
     '''Delete metrics from the given pushgateway.
 
     `gateway` the url for your push gateway. Either of the form
@@ -206,10 +205,10 @@ def delete_from_gateway(gateway, job, grouping_key=None, timeout=None, handler=N
 
     This deletes metrics with the given job and grouping_key.
     This uses the DELETE HTTP method.'''
-    _use_gateway('DELETE', gateway, job, None, grouping_key, timeout, handler, handler_args)
+    _use_gateway('DELETE', gateway, job, None, grouping_key, timeout, handler)
 
 
-def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler, handler_args):
+def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler):
     gateway_url = urlparse(gateway)
     if not gateway_url.scheme:
         gateway = 'http://{0}'.format(gateway)
@@ -227,10 +226,8 @@ def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler,
     headers=[('Content-Type', CONTENT_TYPE_LATEST)]
     if handler is None:
         handler = default_handler
-    if handler_args is None:
-        handler_args = dict()
     handler(url=url, method=method, timeout=timeout,
-            headers=headers, data=data, **handler_args)
+            headers=headers, data=data)()
 
 def instance_ip_grouping_key():
     '''Grouping key with instance set to the IP Address of this host.'''
diff --git a/prometheus_client/handlers/base.py b/prometheus_client/handlers/base.py
@@ -6,13 +6,16 @@
     # Python 3
     from urllib.request import build_opener, Request, HTTPHandler
 
-def handler(url, method, timeout, headers, data, **kwargs):
-    '''Default handler that implements HTTP/HTTPS connections.'''
-    request = Request(url, data=data)
-    request.get_method = lambda: method
-    for k, v in headers:
-        request.add_header(k, v)
-    resp = build_opener(HTTPHandler).open(request, timeout=timeout)
-    if resp.code >= 400:
-        raise IOError("error talking to pushgateway: {0} {1}".format(
-            resp.code, resp.msg))
+def handler(url, method, timeout, headers, data):
+    def handle():
+        '''Default handler that implements HTTP/HTTPS connections.'''
+        request = Request(url, data=data)
+        request.get_method = lambda: method
+        for k, v in headers:
+            request.add_header(k, v)
+        resp = build_opener(HTTPHandler).open(request, timeout=timeout)
+        if resp.code >= 400:
+            raise IOError("error talking to pushgateway: {0} {1}".format(
+                resp.code, resp.msg))
+
+    return handle
diff --git a/prometheus_client/handlers/basic_auth.py b/prometheus_client/handlers/basic_auth.py
@@ -4,13 +4,14 @@
 
 from prometheus_client.handlers.base import handler as default_handler
 
-def handler(url, method, timeout, headers, data, **kwargs):
-    '''Handler that implements HTTP Basic Auth by setting auth headers if
-    'username' passed as keyword argument.'''
-    if 'username' in kwargs:
-        username = kwargs['username']
-        password = kwargs['password']
-        auth_value = "{0}:{1}".format(username, password)
-        auth_header = "Basic {0}".format(base64.b64encode(bytes(auth_value)))
-        headers.append(['Authorization', auth_header])
-    default_handler(url, method, timeout, headers, data)
+def handler(url, method, timeout, headers, data, username = None, password = None):
+    def handle():
+        '''Handler that implements HTTP Basic Auth by setting auth headers if
+        'username' and 'password' arguments are supplied and not None.'''
+        if username is not None and password is not None:
+            auth_value = "{0}:{1}".format(username, password)
+            auth_header = "Basic {0}".format(base64.b64encode(bytes(auth_value)))
+            headers.append(['Authorization', auth_header])
+        default_handler(url, method, timeout, headers, data)()
+
+    return handle
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
@@ -14,6 +14,7 @@
 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.handlers.base import handler as default_handler
+from prometheus_client.handlers.basic_auth import handler as basic_auth_handler
 
 try:
     from BaseHTTPServer import BaseHTTPRequestHandler
@@ -99,7 +100,10 @@ def setUp(self):
         self.requests = requests = []
         class TestHandler(BaseHTTPRequestHandler):
             def do_PUT(self):
-                self.send_response(201)
+                if 'with_basic_auth' in self.requestline and self.headers['authorization'] != 'Basic Zm9vOmJhcg==':
+                    self.send_response(401)
+                else:
+                    self.send_response(201)
                 length = int(self.headers['content-length'])
                 requests.append((self, self.rfile.read(length)))
                 self.end_headers()
@@ -168,14 +172,23 @@ def test_delete_with_groupingkey(self):
     def test_push_with_handler(self):
         def my_test_handler(url, method, timeout, headers, data):
             headers.append(['X-Test-Header', 'foobar'])
-            default_handler(url, method, timeout, headers, data)
+            return default_handler(url, method, timeout, headers, data)
         push_to_gateway(self.address, "my_job", self.registry, handler=my_test_handler)
         self.assertEqual(self.requests[0][0].command, 'PUT')
         self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job')
         self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
         self.assertEqual(self.requests[0][0].headers.get('x-test-header'), 'foobar')
         self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
 
+    def test_push_with_basic_auth_handler(self):
+        def my_auth_handler(url, method, timeout, headers, data):
+            return basic_auth_handler(url, method, timeout, headers, data, "foo", "bar")
+        push_to_gateway(self.address, "my_job_with_basic_auth", self.registry, handler=my_auth_handler)
+        self.assertEqual(self.requests[0][0].command, 'PUT')
+        self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job_with_basic_auth')
+        self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
+        self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
+
     @unittest.skipIf(
         sys.platform == "darwin",
         "instance_ip_grouping_key() does not work on macOS."
