Implement redirect handler option.

prometheus_client/exposition.py

@@ -10,7 +10,7 @@
 
 from .openmetrics import exposition as openmetrics
 from .registry import REGISTRY
-from .utils import floatToGoString
+from .utils import floatToGoString, PrometheusRedirectHandler
 
 try:
     from urllib import quote_plus
@@ -141,7 +141,7 @@ def sample_line(line):
             raise
 
         for suffix, lines in sorted(om_samples.items()):
-            output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix, 
+            output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix,
                                                        metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
             output.append('# TYPE {0}{1} gauge\n'.format(metric.name, suffix))
             output.extend(lines)
@@ -205,24 +205,43 @@ def write_to_textfile(path, registry):
     os.rename(tmppath, path)
 
 
-def default_handler(url, method, timeout, headers, data):
-    """Default handler that implements HTTP/HTTPS connections.
-
-    Used by the push_to_gateway functions. Can be re-used by other handlers."""
+def _make_handler(url, method, timeout, headers, data, base_handler):
 
     def handle():
         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)
+        resp = build_opener(base_handler).open(request, timeout=timeout)
         if resp.code >= 400:
             raise IOError("error talking to pushgateway: {0} {1}".format(
                 resp.code, resp.msg))
 
     return handle
 
 
+def default_handler(url, method, timeout, headers, data):
+    """Default handler that implements HTTP/HTTPS connections.
+
+    Used by the push_to_gateway functions. Can be re-used by other handlers."""
+
+    return _make_handler(url, method, timeout, headers, data, HTTPHandler)
+
+
+def passthrough_redirect_handler(url, method, timeout, headers, data):
+    """
+    Handler that automatically trusts redirect responses for all HTTP methods.
+
+    Augments standard HTTPRedirectHandler capability by permitting PUT requests,
+    preserving the method upon redirect, and passing through all headers and
+    data from the original request. Only use this handler if you control or
+    trust the source of redirect responses you encounter when making requests
+    via the Prometheus client. This handler will simply repeat the identical
+    request, including same method and data, to the new redirect URL."""
+
+    return _make_handler(url, method, timeout, headers, data, PrometheusRedirectHandler)
+
+
 def basic_auth_handler(url, method, timeout, headers, data, username=None, password=None):
     """Handler that implements HTTP/HTTPS connections with Basic Auth.
 

prometheus_client/utils.py

@@ -1,5 +1,14 @@
 import math
+import sys
 
+try:
+    from urllib2 import HTTPError, HTTPRedirectHandler, Request
+except ImportError:
+    # Python 3
+    from urllib.error import HTTPError
+    from urllib.request import HTTPRedirectHandler, Request
+
+PYTHON27_OR_OLDER = sys.version_info <= (2, 7)
 INF = float("inf")
 MINUS_INF = float("-inf")
 NaN = float("NaN")
@@ -22,3 +31,62 @@ def floatToGoString(d):
             mantissa = '{0}.{1}{2}'.format(s[0], s[1:dot], s[dot + 1:]).rstrip('0.')
             return '{0}e+0{1}'.format(mantissa, dot - 1)
         return s
+
+
+class PrometheusRedirectHandler(HTTPRedirectHandler):
+    """
+    Allow additional methods (e.g. PUT) and data forwarding in redirects.
+
+    Use of this class constitute a user's explicit agreement to the
+    redirect responses the Prometheus client will receive when using it.
+    You should only use this class if you control or otherwise trust the
+    redirect behavior involved and are certain it is safe to full transfer
+    the original request (method and data) to the redirected URL. For
+    example, if you know there is a cosmetic URL redirect in front of a
+    local deployment of a Prometheus server, and all redirects are safe,
+    this is the class to use to handle redirects in that case.
+
+    The standard HTTPRedirectHandler does not forward request data nor
+    does it allow redirected PUT requests (which Prometheus uses for some
+    operations, for example `push_to_gateway`) because these cannot
+    generically guarantee no violations of HTTP RFC 2616 requirements for
+    the user to explicitly confirm redirects that could have unexpected
+    side effects (such as rendering a PUT request non-idempotent or
+    creating multiple resources not named in the original request).
+    """
+
+    def redirect_request(self, req, fp, code, msg, headers, newurl):
+        """
+        Apply redirect logic to a request.
+
+        See parent HTTPRedirectHandler.redirect_request for parameter info.
+
+        If the redirect is disallowed, this raises the corresponding HTTP error.
+        If the redirect can't be determined, return None to allow other handlers
+        to try. If the redirect is allowed, return the new request.
+
+        This method specialized for the case when (a) the user knows that the
+        redirect will not cause unacceptable side effects for any request method,
+        and (b) the user knows that any request data should be passed through to
+        the redirect. If either condition is not met, this should not be used.
+        """
+        # note that requests being provided by a handler will use get_method to
+        # indicate the method, by monkeypatching this, instead of setting the
+        # Request object's method attribute.
+        m = getattr(req, "method", req.get_method())
+        if not (code in (301, 302, 303, 307) and m in ("GET", "HEAD")
+                or code in (301, 302, 303) and m in ("POST", "PUT")):
+            raise HTTPError(req.full_url, code, msg, headers, fp)
+        new_request = Request(
+            newurl.replace(' ', '%20'),  # space escaping in new url if needed.
+            headers=req.headers,
+            origin_req_host=req.origin_req_host,
+            unverifiable=True,
+            data=req.data,
+        )
+        if PYTHON27_OR_OLDER:
+            # the `method` attribute did not exist for Request in Python 2.7.
+            new_request.get_method = lambda: m
+        else:
+            new_request.method = m
+        return new_request

tests/test_exposition.py

@@ -14,6 +14,7 @@
 from prometheus_client.core import GaugeHistogramMetricFamily, Timestamp
 from prometheus_client.exposition import (
     basic_auth_handler, default_handler, MetricsHandler,
+    passthrough_redirect_handler,
 )
 
 if sys.version_info < (2, 7):
@@ -208,6 +209,7 @@ def collect(self):
 
 class TestPushGateway(unittest.TestCase):
     def setUp(self):
+        redirect_flag = 'testFlag'
         self.registry = CollectorRegistry()
         self.counter = Gauge('g', 'help', registry=self.registry)
         self.requests = requests = []
@@ -216,6 +218,11 @@ class TestHandler(BaseHTTPRequestHandler):
             def do_PUT(self):
                 if 'with_basic_auth' in self.requestline and self.headers['authorization'] != 'Basic Zm9vOmJhcg==':
                     self.send_response(401)
+                elif 'redirect' in self.requestline and redirect_flag not in self.requestline:
+                    # checks for an initial test request with 'redirect' but without the redirect_flag,
+                    # and simulates a redirect to a url with the redirect_flag (which will produce a 201)
+                    self.send_response(301)
+                    self.send_header('Location', getattr(self, 'redirect_address', None))
                 else:
                     self.send_response(201)
                 length = int(self.headers['content-length'])
@@ -225,6 +232,22 @@ def do_PUT(self):
             do_POST = do_PUT
             do_DELETE = do_PUT
 
+        # set up a separate server to serve a fake redirected request.
+        # the redirected URL will have `redirect_flag` added to it,
+        # which will cause the request handler to return 201.
+        httpd_redirect = HTTPServer(('localhost', 0), TestHandler)
+        self.redirect_address = TestHandler.redirect_address = \
+            'http://localhost:{0}/{1}'.format(httpd_redirect.server_address[1], redirect_flag)
+
+        class TestRedirectServer(threading.Thread):
+            def run(self):
+                httpd_redirect.handle_request()
+
+        self.redirect_server = TestRedirectServer()
+        self.redirect_server.daemon = True
+        self.redirect_server.start()
+
+        # set up the normal server to serve the example requests across test cases.
         httpd = HTTPServer(('localhost', 0), TestHandler)
         self.address = 'http://localhost:{0}'.format(httpd.server_address[1])
 
@@ -236,6 +259,7 @@ def run(self):
         self.server.daemon = True
         self.server.start()
 
+
     def test_push(self):
         push_to_gateway(self.address, "my_job", self.registry)
         self.assertEqual(self.requests[0][0].command, 'PUT')
@@ -330,6 +354,16 @@ def my_auth_handler(url, method, timeout, headers, data):
         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')
 
+    def test_push_with_redirect_handler(self):
+        def my_redirect_handler(url, method, timeout, headers, data):
+            return passthrough_redirect_handler(url, method, timeout, headers, data)
+
+        push_to_gateway(self.address, "my_job_with_redirect", self.registry, handler=my_redirect_handler)
+        self.assertEqual(self.requests[0][0].command, 'PUT')
+        self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job_with_redirect')
+        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."