8000 Merge pull request #126 from rossigee/master · dds/client_python@aa4b4fb · GitHub
[go: up one dir, main page]
Skip to content

Commit aa4b4fb

Browse files
brian-brazilbrian-brazil
authored
Merge pull request prometheus#126 from rossigee/master
* Allow a handler to be passed in to carry out a custom request. 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.
2 parents fc6aa20 + a7afc96 commit aa4b4fb

File tree

3 files changed

+133
-18
lines changed

3 files changed

+133
-18
lines changed

‎README.md

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,24 @@ for more information.
321321
`instance_ip_grouping_key` returns a grouping key with the instance label set
322322
to the host's IP address.
323323

324+
### Handlers for authentication
325+
326+
If the push gateway you are connecting to is protected with HTTP Basic Auth,
327+
you can use a special handler to set the Authorization header.
328+
329+
```python
330+
from prometheus_client import CollectorRegistry, Gauge, push_to_gateway
331+
from prometheus_client.exposition import basic_auth_handler
332+
333+
def my_auth_handler(url, method, timeout, headers, data):
334+
username = 'foobar'
335+
password = 'secret123'
336+
return basic_auth_handler(url, method, timeout, headers, data, username, password)
337+
registry = CollectorRegistry()
338+
g = Gauge('job_last_success_unixtime', 'Last time a batch job successfully finished', registry=registry)
339+
g.set_to_current_time()
340+
push_to_gateway('localhost:9091', job='batchA', registry=registry, handler=my_auth_handler)
341+
```
324342

325343
## Bridges
326344

‎prometheus_client/exposition.py

Lines changed: 89 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import threading
99
from contextlib import closing
1010
from wsgiref.simple_server import make_server
11+
import base64
12+
import sys
1113

1214
from . import core
1315
try:
@@ -118,7 +120,46 @@ def write_to_textfile(path, registry):
118120
os.rename(tmppath, path)
119121

120122

121-
def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None):
123+
def default_handler(url, method, timeout, headers, data):
124+
'''Default handler that implements HTTP/HTTPS connections.
125+
126+
Used by the push_to_gateway functions. Can be re-used by other handlers.'''
127+
def handle():
128+
request = Request(url, data=data)
129+
request.get_method = lambda: method
130+
for k, v in headers:
131+
request.add_header(k, v)
132+
resp = build_opener(HTTPHandler).open(request, timeout=timeout)
133+
if resp.code >= 400:
134+
raise IOError("error talking to pushgateway: {0} {1}".format(
135+
resp.code, resp.msg))
136+
137+
return handle
138+
139+
140+
def basic_auth_handler(url, method, timeout, headers, data, username=None, password=None):
141+
'''Handler that implements HTTP/HTTPS connections with Basic Auth.
142+
143+
Sets auth headers using supplied 'username' and 'password', if set.
144+
Used by the push_to_gateway functions. Can be re-used by other handlers.'''
145+
def handle():
146+
'''Handler that implements HTTP Basic Auth.
147+
'''
148+
if username is not None and password is not None:
149+
if sys.version_info >= (3,0):
150+
auth_value = bytes('{0}:{1}'.format(username, password), 'utf8')
151+
auth_token = str(base64.b64encode(auth_value), 'utf8')
152+
else:
153+
auth_value = '{0}:{1}'.format(username, password)
154+
auth_token = base64.b64encode(auth_value)
155+
auth_header = "Basic {0}".format(auth_token)
156+
headers.append(['Authorization', auth_header])
157+
default_handler(url, method, timeout, headers, data)()
158+
159+
return handle
160+
161+
162+
def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=default_handler):
122163
'''Push metrics to the given pushgateway.
123164
124165
`gateway` the url for your push gateway. Either of the form
@@ -130,13 +171,37 @@ def push_to_gateway(gateway, job, registry, grouping_key=None, timeout=None):
130171
Defaults to None
131172
`timeout` is how long push will attempt to connect before giving up.
132173
Defaults to None
174+
`handler` is an optional function which can be provided to perform
175+
requests to the 'gateway'.
176+
Defaults to None, in which case an http or https request
177+
will be carried out by a default handler.
178+
If not None, the argument must be a function which accepts
179+
the following arguments:
180+
url, method, timeout, headers, and content
181+
May be used to implement additional functionality not
182+
supported by the built-in default handler (such as SSL
183+
client certicates, and HTTP authentication mechanisms).
184+
'url' is the URL for the request, the 'gateway' argument
185+
described earlier will form the basis of this URL.
186+
'method' is the HTTP method which should be used when
187+
carrying out the request.
188+
'timeout' requests not successfully completed after this
189+
many seconds should be aborted. If timeout is None, then
190+
the handler should not set a timeout.
191+
F438 'headers' is a list of ("header-name","header-value") tuples
192+
which must be passed to the pushgateway in the form of HTTP
193+
request headers.
194+
The function should raise an exception (e.g. IOError) on
195+
failure.
196+
'content' is the data which should be used to form the HTTP
197+
Message Body.
133198
134199
This overwrites all metrics with the same job and grouping_key.
135200
This uses the PUT HTTP method.'''
136-
_use_gateway('PUT', gateway, job, registry, grouping_key, timeout)
201+
_use_gateway('PUT', gateway, job, registry, grouping_key, timeout, handler)
137202

138203

139-
def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None):
204+
def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None, handler=default_handler):
140205
'''PushAdd metrics to the given pushgateway.
141206
142207
`gateway` the url for your push gateway. Either of the form
@@ -148,13 +213,19 @@ def pushadd_to_gateway(gateway, job, registry, grouping_key=None, timeout=None):
148213
Defaults to None
149214
`timeout` is how long push will attempt to connect before giving up.
150215
Defaults to None
216+
`handler` is an optional function which can be provided to perform
217+
requests to the 'gateway'.
218+
Defaults to None, in which case an http or https request
219+
will be carried out by a default handler.
220+
See the 'prometheus_client.push_to_gateway' documentation
221+
for implementation requirements.
151222
152223
This replaces metrics with the same name, job and grouping_key.
153224
This uses the POST HTTP method.'''
154-
_use_gateway('POST', gateway, job, registry, grouping_key, timeout)
225+
_use_gateway('POST', gateway, job, registry, grouping_key, timeout, handler)
155226

156227

157-
def delete_from_gateway(gateway, job, grouping_key=None, timeout=None):
228+
def delete_from_gateway(gateway, job, grouping_key=None, timeout=None, handler=default_handler):
158229
'''Delete metrics from the given pushgateway.
159230
160231
`gateway` the url for your push gateway. Either of the form
@@ -165,14 +236,21 @@ def delete_from_gateway(gateway, job, grouping_key=None, timeout=None):
165236
Defaults to None
166237
`timeout` is how long delete will attempt to connect before giving up.
167238
Defaults to None
239+
`handler` is an optional function which can be provided to perform
240+
requests to the 'gateway'.
241+
Defaults to None, in which case an http or https request
242+
will be carried out by a default handler.
243+
See the 'prometheus_client.push_to_gateway' documentation
244+
for implementation requirements.
168245
169246
This deletes metrics with the given job and grouping_key.
170247
This uses the DELETE HTTP method.'''
171-
_use_gateway('DELETE', gateway, job, None, grouping_key, timeout)
248+
_use_gateway('DELETE', gateway, job, None, grouping_key, timeout, handler)
172249

173250

174-
def _use_gateway(method, gateway, job, registry, grouping_key, timeout):
175-
if not (gateway.startswith('http://') or gateway.startswith('https://')):
251+
def _use_gateway(method, gateway, job, registry, grouping_key, timeout, handler):
252+
gateway_url = urlparse(gateway)
253+
if not gateway_url.scheme:
176254
gateway = 'http://{0}'.format(gateway)
177255
url = '{0}/metrics/job/{1}'.format(gateway, quote_plus(job))
178256

@@ -185,13 +263,9 @@ def _use_gateway(method, gateway, job, registry, grouping_key, timeout):
185263
url = url + ''.join(['/{0}/{1}'.format(quote_plus(str(k)), quote_plus(str(v)))
186264
for k, v in sorted(grouping_key.items())])
187265

188-
request = Request(url, data=data)
189-
request.add_header('Content-Type', CONTENT_TYPE_LATEST)
190-
request.get_method = lambda: method
191-
resp = build_opener(HTTPHandler).open(request, timeout=timeout)
192-
if resp.code >= 400:
193-
raise IOError("error talking to pushgateway: {0} {1}".format(
194-
resp.code, resp.msg))
266+
headers=[('Content-Type', CONTENT_TYPE_LATEST)]
267+
handler(url=url, method=method, timeout=timeout,
268+
headers=headers, data=data)()
195269

196270
def instance_ip_grouping_key():
197271
'''Grouping key with instance set to the IP Address of this host.'''

‎tests/test_exposition.py

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from prometheus_client import CollectorRegistry, generate_latest
1414
from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
1515
from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
16+
from prometheus_client.exposition import default_handler, basic_auth_handler
1617

1718
try:
1819
from BaseHTTPServer import BaseHTTPRequestHandler
@@ -22,7 +23,6 @@
2223
from http.server import BaseHTTPRequestHandler
2324
from http.server import HTTPServer
2425

25-
2626
class TestGenerateText(unittest.TestCase):
2727
def setUp(self):
2828
self.registry = CollectorRegistry()
@@ -99,7 +99,10 @@ def setUp(self):
9999
self.requests = requests = []
100100
class TestHandler(BaseHTTPRequestHandler):
101101
def do_PUT(self):
102-
self.send_response(201)
102+
if 'with_basic_auth' in self.requestline and self.headers['authorization'] != 'Basic Zm9vOmJhcg==':
103+
self.send_response(401)
104+
else:
105+
self.send_response(201)
103106
length = int(self.headers['content-length'])
104107
requests.append((self, self.rfile.read(length)))
105108
self.end_headers()
@@ -108,7 +111,7 @@ def do_PUT(self):
108111
do_DELETE = do_PUT
109112

110113
httpd = HTTPServer(('localhost', 0), TestHandler)
111-
self.address = ':'.join([str(x) for x in httpd.server_address])
114+
self.address = 'http://localhost:{0}'.format(httpd.server_address[1])
112115
class TestServer(threading.Thread):
113116
def run(self):
114117
httpd.handle_request()
@@ -165,6 +168,26 @@ def test_delete_with_groupingkey(self):
165168
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
166169
self.assertEqual(self.requests[0][1], b'')
167170

171+
def test_push_with_handler(self):
172+
def my_test_handler(url, method, timeout, headers, data):
173+
headers.append(['X-Test-Header', 'foobar'])
174+
return default_handler(url, method, timeout, headers, data)
175+
push_to_gateway(self.address, "my_job", self.registry, handler=my_test_handler)
176+
self.assertEqual(self.requests[0][0].command, 'PUT')
177+
self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job')
178+
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
179+
self.assertEqual(self.requests[0][0].headers.get('x-test-header'), 'foobar')
180+
self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
181+
182+
def test_push_with_basic_auth_handler(self):
183+
def my_auth_handler(url, method, timeout, headers, data):
184+
return basic_auth_handler(url, method, timeout, headers, data, "foo", "bar")
185+
push_to_gateway(self.address, "my_job_with_basic_auth", self.registry, handler=my_auth_handler)
186+
self.assertEqual(self.requests[0][0].command, 'PUT')
187+
self.assertEqual(self.requests[0][0].path, '/metrics/job/my_job_with_basic_auth')
188+
self.assertEqual(self.requests[0][0].headers.get('content-type'), CONTENT_TYPE_LATEST)
189+
self.assertEqual(self.requests[0][1], b'# HELP g help\n# TYPE g gauge\ng 0.0\n')
190+
168191
@unittest.skipIf(
169192
sys.platform == "darwin",
170193
"instance_ip_grouping_key() does not work on macOS."

0 commit comments

Comments
 (0)
0