Skip to content

Commit

Permalink
feat: add proxy_headers to configuration options
Browse files Browse the repository at this point in the history
  • Loading branch information
bednar committed Aug 12, 2021
1 parent 2db92d4 commit 020a6cc
Show file tree
Hide file tree
Showing 5 changed files with 76 additions and 1 deletion.
3 changes: 3 additions & 0 deletions influxdb_client/client/influxdb_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,8 @@ def __init__(self, url, token, debug=None, timeout=10_000, enable_gzip=False, or
:key bool verify_ssl: Set this to false to skip verifying SSL certificate when calling API from https server.
:key str ssl_ca_cert: Set this to customize the certificate file to verify the peer.
:key str proxy: Set this to configure the http proxy to be used (ex. http://localhost:3128)
:key str proxy_headers: A dictionary containing headers that will be sent to the proxy. Could be used for proxy
authentication.
:key int connection_pool_maxsize: Number of connections to save that can be reused by urllib3.
Defaults to "multiprocessing.cpu_count() * 5".
:key urllib3.util.retry.Retry retries: Set the default retry strategy that is used for all HTTP requests
Expand All @@ -63,6 +65,7 @@ def __init__(self, url, token, debug=None, timeout=10_000, enable_gzip=False, or
conf.verify_ssl = kwargs.get('verify_ssl', True)
conf.ssl_ca_cert = kwargs.get('ssl_ca_cert', None)
conf.proxy = kwargs.get('proxy', None)
conf.proxy_headers = kwargs.get('proxy_headers', None)
conf.connection_pool_maxsize = kwargs.get('connection_pool_maxsize', conf.connection_pool_maxsize)
conf.timeout = timeout

Expand Down
2 changes: 2 additions & 0 deletions influxdb_client/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,8 @@ def __init__(self):

# Proxy URL
self.proxy = None
# A dictionary containing headers that will be sent to the proxy
self.proxy_headers = None
# Safe chars for path_param
self.safe_chars_for_path_param = ''

Expand Down
1 change: 1 addition & 0 deletions influxdb_client/rest.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,7 @@ def __init__(self, configuration, pools_size=4, maxsize=None, retries=False):
cert_file=configuration.cert_file,
key_file=configuration.key_file,
proxy_url=configuration.proxy,
proxy_headers=configuration.proxy_headers,
**addition_pool_args
)
else:
Expand Down
47 changes: 46 additions & 1 deletion tests/test_InfluxDBClient.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
import unittest

from influxdb_client import InfluxDBClient, Point
from influxdb_client.client.write_api import SYNCHRONOUS, ASYNCHRONOUS, WriteOptions, WriteType
from influxdb_client.client.write_api import WriteOptions, WriteType

from tests.base_test import BaseTest


class InfluxDBClientTest(unittest.TestCase):
Expand Down Expand Up @@ -167,6 +169,49 @@ def test_write_context_manager(self):
self.assertIsNone(api_client._pool)
self.assertIsNone(self.client.api_client)


class InfluxDBClientTestIT(BaseTest):
httpRequest = []

def tearDown(self) -> None:
super(InfluxDBClientTestIT, self).tearDown()
if hasattr(self, 'httpd'):
self.httpd.shutdown()
if hasattr(self, 'httpd_thread'):
self.httpd_thread.join()
InfluxDBClientTestIT.httpRequest = []

def test_proxy(self):
self._start_proxy_server()

self.client.close()
self.client = InfluxDBClient(url=self.host,
token=self.auth_token,
proxy=f"http://localhost:{self.httpd.server_address[1]}",
proxy_headers={'ProxyHeader': 'Val'})
ready = self.client.ready()
self.assertEqual(ready.status, "ready")
self.assertEqual(1, len(InfluxDBClientTestIT.httpRequest))
self.assertEqual('Val', InfluxDBClientTestIT.httpRequest[0].headers.get('ProxyHeader'))

def _start_proxy_server(self):
import http.server
import urllib.request

class ProxyHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):

def do_GET(self):
InfluxDBClientTestIT.httpRequest.append(self)
self.send_response(200)
self.send_header('Content-type', 'application/json')
self.end_headers()
self.copyfile(urllib.request.urlopen(self.path), self.wfile)

self.httpd = http.server.HTTPServer(('localhost', 0), ProxyHTTPRequestHandler)
self.httpd_thread = threading.Thread(target=self.httpd.serve_forever)
self.httpd_thread.start()


class ServerWithSelfSingedSSL(http.server.SimpleHTTPRequestHandler):
def _set_headers(self):
self.send_response(200)
Expand Down
24 changes: 24 additions & 0 deletions tests/test_WriteApi.py
Original file line number Diff line number Diff line change
Expand Up @@ -515,6 +515,30 @@ def test_writes_default_tags_dict_without_tag(self):
self.assertEqual(1, len(requests))
self.assertEqual("h2o,customer=California\\ Miner,id=132-987-655 level=1 1", requests[0].parsed_body)

def test_redirect(self):
from urllib3 import Retry
Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT = frozenset()
Retry.DEFAULT.remove_headers_on_redirect = Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT
self.influxdb_client.close()

self.influxdb_client = InfluxDBClient(url="http://localhost", token="my-token", org="my-org")

httpretty.register_uri(httpretty.POST, uri="http://localhost2/api/v2/write", status=204)
httpretty.register_uri(httpretty.POST, uri="http://localhost/api/v2/write", status=301,
adding_headers={'Location': 'http://localhost2/api/v2/write'})

self.write_client = self.influxdb_client.write_api(write_options=SYNCHRONOUS)

self.write_client.write("my-bucket", "my-org", {"measurement": "h2o", "fields": {"level": 1.0}, "time": 1})

requests = httpretty.httpretty.latest_requests
self.assertEqual(2, len(requests))
self.assertEqual('Token my-token', requests[0].headers['Authorization'])
self.assertEqual('Token my-token', requests[1].headers['Authorization'])

from urllib3 import Retry
Retry.DEFAULT.remove_headers_on_redirect = Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT


class AsynchronousWriteTest(BaseTest):

Expand Down

0 comments on commit 020a6cc

Please sign in to comment.