paraizofelipe
diff --git a/‎prometheus_client/exposition.py
Lines changed: 92 additions & 9 deletions b/‎prometheus_client/exposition.py
Lines changed: 92 additions & 9 deletions
diff --git a/‎tests/test_exposition.py
Lines changed: 46 additions & 0 deletions b/‎tests/test_exposition.py
Lines changed: 46 additions & 0 deletions
@@ -17,22 +17,86 @@
 
     from BaseHTTPServer import BaseHTTPRequestHandler
     from SocketServer import ThreadingMixIn
-    from urllib2 import build_opener, HTTPHandler, Request
+    from urllib2 import (
+        build_opener, HTTPError, HTTPHandler, HTTPRedirectHandler, Request,
+    )
     from urlparse import parse_qs, urlparse
 except ImportError:
     # Python 3
     from http.server import BaseHTTPRequestHandler
     from socketserver import ThreadingMixIn
+    from urllib.error import HTTPError
     from urllib.parse import parse_qs, quote_plus, urlparse
-    from urllib.request import build_opener, HTTPHandler, Request
+    from urllib.request import (
+        build_opener, HTTPHandler, HTTPRedirectHandler, Request,
+    )
 
 CONTENT_TYPE_LATEST = str('text/plain; version=0.0.4; charset=utf-8')
 """Content type of the latest text format"""
-
+PYTHON27_OR_OLDER = sys.version_info < (3, )
 PYTHON26_OR_OLDER = sys.version_info < (2, 7)
 PYTHON376_OR_NEWER = sys.version_info > (3, 7, 5)
 
 
+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
+
+
 def _bake_output(registry, accept_header, params):
     """Bake output for metrics output."""
     encoder, content_type = choose_encoder(accept_header)
@@ -141,7 +205,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 +269,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.
 
 
@@ -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,8 @@ def collect(self):
 
 class TestPushGateway(unittest.TestCase):
     def setUp(self):
+        redirect_flag = 'testFlag'
+        self.redirect_flag = redirect_flag  # preserve a copy for downstream test assertions
         self.registry = CollectorRegistry()
         self.counter = Gauge('g', 'help', registry=self.registry)
         self.requests = requests = []
@@ -216,6 +219,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 +233,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 +260,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 +355,27 @@ 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')
+
+        # ensure the redirect preserved request settings from the initial request.
+        self.assertEqual(self.requests[0][0].command, self.requests[1][0].command)
+        self.assertEqual(
+            self.requests[0][0].headers.get('content-type'),
+            self.requests[1][0].headers.get('content-type')
+        )
+        self.assertEqual(self.requests[0][1], self.requests[1][1])
+
+        # ensure the redirect took place at the expected redirect location.
+        self.assertEqual(self.requests[1][0].path, "/" + self.redirect_flag)
+
     @unittest.skipIf(
         sys.platform == "darwin",
         "instance_ip_grouping_key() does not work on macOS."