feat: add proxy_headers to configuration options

Author: bednar

influxdb_client/client/influxdb_client.py

@@ -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
@@ -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
 

influxdb_client/configuration.py

@@ -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 = ''
 

influxdb_client/rest.py

@@ -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:

tests/test_InfluxDBClient.py

@@ -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):
@@ -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)

tests/test_WriteApi.py

@@ -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):