Allow a handler to be passed in to carry out a custom request.

tim-seoss · rossigee · commit c3de82b60c11 · 2017-01-25T10:35:47.000+07:00
Allow a custom handler to be provided, so that the caller can provide
code which carried out basic auth, https client certificate validation
or other arbitrary schemes and access methods such as using different
types of proxy.
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -118,7 +118,7 @@ def write_to_textfile(path, registry):
     os.rename(tmppath, path)
 
 
-def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=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
@@ -133,10 +133,10 @@ def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None):
 
     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)
+    _use_gateway('PUT', gateway, job, registry, grouping_key, timeout, handler)
 
 
-def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=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
@@ -151,10 +151,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)
+    _use_gateway('POST', gateway, job, registry, grouping_key, timeout, handler)
 
 
-def delete_from_gateway(gateway, job, grouping_key=None, timeout=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
@@ -168,10 +168,10 @@ def delete_from_gateway(gateway, job, grouping_key=None, timeout=None):
 
     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)
+    _use_gateway('DELETE', gateway, job, None, grouping_key, timeout, handler)
 
 
-def _use_gateway(method, gateway, job, registry, grouping_key, timeout):
+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)
@@ -189,10 +189,14 @@ def _use_gateway(method, gateway, job, registry, grouping_key, timeout):
     request = Request(url, data=data)
     request.add_header('Content-Type', CONTENT_TYPE_LATEST)
     request.get_method = lambda: method
-    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))
+    if handler is None:
+        resp = build_opener(handler).open(request, timeout=timeout)
+        if resp.code >= 400:
+            raise IOError("error talking to pushgateway: {0} {1}".format(
+                resp.code, resp.msg))
+    else:
+        handler(url=url, method=lambda: method, timeout=timeout,
+                headers=[('Content-Type', CONTENT_TYPE_LATEST)], content=data)
 
 def instance_ip_grouping_key():
     '''Grouping key with instance set to the IP Address of this host.'''
