8000 Implement redirect handler option. (#622) · ethervoid/client_python@5e3674a · GitHub
[go: up one dir, main page]
Skip to content

Commit 5e3674a

Browse files
espears1espears1
authored
Implement redirect handler option. (prometheus#622)
* Implement redirect handler option. * fix isort sort order * Make redirect handler compatible with Python 2.7. * Move handler class to private class in exposition.py * Add more redirect checks to unit tests. * Fix isort formatting. * Fix Python 2.7 version check Signed-off-by: Ely Spears <espears@squarespace.com>
1 parent ccb8395 commit 5e3674a

File tree

2 files changed

+138
-9
lines changed

2 files changed

+138
-9
lines changed

‎prometheus_client/exposition.py

Lines changed: 92 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,22 +17,86 @@
1717

1818
from BaseHTTPServer import BaseHTTPRequestHandler
1919
from SocketServer import ThreadingMixIn
20-
from urllib2 import build_opener, HTTPHandler, Request
20+
from urllib2 import (
21+
build_opener, HTTPError, HTTPHandler, HTTPRedirectHandler, Request,
22+
)
2123
from urlparse import parse_qs, urlparse
2224
except ImportError:
2325
# Python 3
2426
from http.server import BaseHTTPRequestHandler
2527
from socketserver import ThreadingMixIn
28+
from urllib.error import HTTPError
2629
from urllib.parse import parse_qs, quote_plus, urlparse
27-
from urllib.request import build_opener, HTTPHandler, Request
30+
from urllib.request import (
31+
build_opener, HTTPHandler, HTTPRedirectHandler, Request,
32+
)
2833

2934
CONTENT_TYPE_LATEST = str('text/plain; version=0.0.4; charset=utf-8')
3035
"""Content type of the latest text format"""
31-
36+
PYTHON27_OR_OLDER = sys.version_info < (3, )
3237
PYTHON26_OR_OLDER = sys.version_info < (2, 7)
3338
PYTHON376_OR_NEWER = sys.version_info > (3, 7, 5)
3439

3540

41+
class _PrometheusRedirectHandler(HTTPRedirectHandler):
42+
"""
43+
Allow additional methods (e.g. PUT) and data forwarding in redirects.
44+
45+
Use of this class constitute a user's explicit agreement to the
46+
redirect responses the Prometheus client will receive when using it.
47+
You should only use this class if you control or otherwise trust the
48+
redirect behavior involved and are certain it is safe to full transfer
49+
the original request (method and data) to the redirected URL. For
50+
example, if you know there is a cosmetic URL redirect in front of a
51+
local deployment of a Prometheus server, and all redirects are safe,
52+
this is the class to use to handle redirects in that case.
53+
54+
The standard HTTPRedirectHandler does not forward request data nor
55+
does it allow redirected PUT requests (which Prometheus uses for some
56+
operations, for example `push_to_gateway`) because these cannot
57+
generically guarantee no violations of HTTP RFC 2616 requirements for
58+
the user to explicitly confirm redirects that could have unexpected
59+
side effects (such as rendering a PUT request non-idempotent or
60+
creating multiple resources not named in the original request).
61+
"""
62+
63+
def redirect_request(self, req, fp, code, msg, headers, newurl):
64+
"""
65+
Apply redirect logic to a request.
66+
67+
See parent HTTPRedirectHandler.redirect_request for parameter info.
68+
69+
If the redirect is disallowed, this raises the corresponding HTTP error.
70+
If the redirect can't be determined, return None to allow other handlers
71+
to try. If the redirect is allowed, return the new request.
72+
73+
This method specialized for the case when (a) the user knows that the
74+
redirect will not cause unacceptable side effects for any request method,
75+
and (b) the user knows that any request data should be passed through to
76+
the redirect. If either condition is not met, this should not be used.
77+
"""
78+
# note that requests being provided by a handler will use get_method to
79+
# indicate the method, by monkeypatching this, instead of setting the
80+
# Request object's method attribute.
81+
m = getattr(req, "method", req.get_method())
82+
if not (code in (301, 302, 303, 307) and m in ("GET", "HEAD")
83+
or code in (301, 302, 303) and m in ("POST", "PUT")):
84+
raise HTTPError(req.full_url, code, msg, headers, fp)
85+
new_request = Request(
86+
newurl.replace(' ', '%20'), # space escaping in new url if needed.
87+
headers=req.headers,
88+
origin_req_host=req.origin_req_host,
89+
unverifiable=True,
90+
data=req.data,
91+
)
92+
if PYTHON27_OR_OLDER:
93+
# the `method` attribute did not exist for Request in Python 2.7.
94+
new_request.get_method = lambda: m
95+
else:
96+
new_request.method = m
97+
return new_request
98+
99+
36100
def _bake_output(registry, accept_header, params):
37101
"""Bake output for metrics output."""
38102
encoder, content_type = choose_encoder(accept_header)
@@ -141,7 +205,7 @@ def sample_line(line):
141205
raise
142206

143207
for suffix, lines in sorted(om_samples.items()):
144-
output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix,
208+
output.append('# HELP {0}{1} {2}\n'.format(metric.name, suffix,
145209
metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
146210
output.append('# TYPE {0}{1} gauge\n'.format(metric.name, suffix))
147211
output.extend(lines)
@@ -205,24 +269,43 @@ def write_to_textfile(path, registry):
205269
os.rename(tmppath, path)
206270

207271

208-
def default_handler(url, method, timeout, headers, data):
209-
"""Default handler that implements HTTP/HTTPS connections.
210-
211-
Used by the push_to_gateway functions. Can be re-used by other handlers."""
272+
def _make_handler(url, method, timeout, headers, data, base_handler):
212273

213274
def handle():
214275
request = Request(url, data=data)
215276
request.get_method = lambda: method
216277
for k, v in headers:
217278
request.add_header(k, v)
218-
resp = build_opener(HTTPHandler).open(request, timeout=timeout)
279+
resp = build_opener(base_handler).open(request, timeout=timeout)
219280
if resp.code >= 400:
220281
raise IOError("error talking to pushgateway: {0} {1}".format(
221282
resp.code, resp.msg))
222283

223284
return handle
224285

225286

287+
def default_handler(url, method, timeout, headers, data):
288+
"""Default handler that implements HTTP/HTTPS connections.
289+
290+
Used by the push_to_gateway functions. Can be re-used by other handlers."""
291+
292+
return _make_handler(url, method, timeout, headers, data, HTTPHandler)
293+
294+
295+
def passthrough_redirect_handler(url, method, timeout, headers, data):
296+
"""
297+
Handler that automatically trusts redirect responses for all HTTP methods.
298+
299+
Augments standard HTTPRedirectHandler capability by permitting PUT requests,
300+
preserving the method upon redirect, and passing through all headers and
301+
data from the original request. Only use this handler if you control or
302+
trust the source of redirect responses you encounter when making requests
303+
via the Prometheus client. This handler will simply repeat the identical
304+
request, including same method and data, to the new redirect URL."""
305+
306+
return _make_handler(url, method, timeout, headers, data, _PrometheusRedirectHandler)
307+
308+
226309
def basic_auth_handler(url, method, timeout, headers, data, username=None, password=None):
227310
"""Handler that implements HTTP/HTTPS connections with Basic Auth.
228311

‎tests/test_exposition.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
from prometheus_client.core import GaugeHistogramMetricFamily, Timestamp
1515
from prometheus_client.exposition import (
1616
basic_auth_handler, default_handler, MetricsHandler,
17+
passthrough_redirect_handler,
1718
)
1819

1920
if sys.version_info < (2, 7):
@@ -208,6 +209,8 @@ def collect(self):
208209

209210
class TestPushGateway(unittest.TestCase):
210211
def setUp(self):
212+
redirect_flag = 'testFlag'
213+
self.redirect_flag = redirect_flag # preserve a copy for downstream test assertions
211214
self.registry = CollectorRegistry()
212215
self.counter = Gauge('g', 'help', registry=self.registry)
213216
self.requests = requests = []
@@ -216,6 +219,11 @@ class TestHandler(BaseHTTPRequestHandler):
216219
def do_PUT(self):
217220
if 'with_basic_auth' in self.requestline and self.headers['authorization'] != 'Basic Zm9vOmJhcg==':
218221
self.send_response(401)
222+
elif 'redirect' in self.requestline and redirect_flag not in self.requestline:
223+
# checks for an initial test request with 'redirect' but without the redirect_flag,
224+
# and simulates a redirect to a url with the redirect_flag (which will produce a 201)
225+
self.send_response(301)
226+
self.send_header('Location', getattr(self, 'redirect_address', None))
219227
else:
220228
self.send_response(201)
221229
length = int(self.headers['content-length'])
@@ -225,6 +233,22 @@ def do_PUT(self):
225233
do_POST = do_PUT
226234
do_DELETE = do_PUT
227235

236+
# set up a separate server to serve a fake redirected request.
237+
# the redirected URL will have `redirect_flag` added to it,
238+
# which will cause the request handler to return 201.
239+
httpd_redirect = HTTPServer(('localhost', 0), TestHandler)
240+
self.redirect_address = TestHandler.redirect_address = \
241+
'http://localhost:{0}/{1}'.format(httpd_redirect.server_address[1], redirect_flag)
242+
243+
class TestRedirectServer(threading.Thread):
244+
def run(self):
245+
httpd_redirect.handle_request()
246+
247+
self.redirect_server = TestRedirectServer()
248+
self.redirect_server.daemon = True
249+
self.redirect_server.start()
250+
251+
# set up the normal server to serve the example requests across test cases.
228252
httpd = HTTPServer(('localhost', 0), TestHandler)
229253
self.address = 'http://localhost:{0}'.format(httpd.server_address[1])
230254

@@ -236,6 +260,7 @@ def run(self):
236260
self.server.daemon = True
237261
self.server.start()
238262

263+
239264
def test_push(self):
240265
push_to_gateway(self.address, "my_job", self.registry)
241266
self.assertEqual(self.requests[0][0].command, 'PUT')
@@ -330,6 +355,27 @@ def my_auth_handler(url, method, timeout, headers, data):
330355
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
331356
self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
332357

358+
def test_push_with_redirect_handler(self):
359+
def my_redirect_handler(url, method, timeout, headers, data):
360+
return passthrough_redirect_handler(url, method, timeout, headers, data)
361+
362+
push_to_gateway(self.address, "my_job_with_redirect", self.registry, handler=my_redirect_handler)
363+
self.assertEqual(self.requests[0][0].command, 'PUT')
364+
self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job_with_redirect')
365+
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
366+
self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
367+
368+
# ensure the redirect preserved request settings from the initial request.
369+
self.assertEqual(self.requests[0][0].command, self.requests[1][0].command)
370+
self.assertEqual(
371+
self.requests[0][0].headers.get('content-type'),
372+
self.requests[1][0].headers.get('content-type')
373+
)
374+
self.assertEqual(self.requests[0][1], self.requests[1][1])
375+
376+
# ensure the redirect took place at the expected redirect location.
377+
self.assertEqual(self.requests[1][0].path, "/" + self.redirect_flag)
378+
333379
@unittest.skipIf(
334380
sys.platform == "darwin",
335381
"instance_ip_grouping_key() does not work on macOS."

0 commit comments

Comments
 (0)
0