feat: add proxy_headers to configuration options (#301)

Author: bednar

CHANGELOG.md

@@ -3,6 +3,10 @@
 ### Features
 1. [#281](https://github.com/influxdata/influxdb-client-python/pull/281): `FluxTable`, `FluxColumn` and `FluxRecord` objects have helpful reprs
 1. [#293](https://github.com/influxdata/influxdb-client-python/pull/293): `dataframe_serializer` supports batching
+1. [#301](https://github.com/influxdata/influxdb-client-python/pull/301): Add `proxy_headers` to configuration options
+
+### Documentation
+1. [#301](https://github.com/influxdata/influxdb-client-python/pull/301): How to configure proxy
 
 ### Bug Fixes
 1. [#283](https://github.com/influxdata/influxdb-client-python/pull/283): Set proxy server in config file

README.rst

@@ -76,6 +76,7 @@ InfluxDB 2.0 client features
     - `How to use Jupyter + Pandas + InfluxDB 2`_
 - Advanced Usage
     - `Gzip support`_
+    - `Proxy configuration`_
     - `Delete data`_
 
 Installation
@@ -1059,6 +1060,41 @@ Gzip support
 
 .. marker-gzip-end
 
+Proxy configuration
+^^^^^^^^^^^^^^^^^^^
+.. marker-proxy-start
+
+You can configure the client to tunnel requests through an HTTP proxy.
+The following proxy options are supported:
+
+- ``proxy`` - Set this to configure the http proxy to be used, ex. ``http://localhost:3128``
+- ``proxy_headers`` - A dictionary containing headers that will be sent to the proxy. Could be used for proxy authentication.
+
+.. code-block:: python
+
+   from influxdb_client import InfluxDBClient
+
+   with InfluxDBClient(url="http://localhost:8086",
+                       token="my-token",
+                       org="my-org",
+                       proxy="http://localhost:3128") as client:
+
+.. note::
+
+    If your proxy notify the client with permanent redirect (``HTTP 301``) to **different host**.
+    The client removes ``Authorization`` header, because otherwise the contents of ``Authorization`` is sent to third parties
+    which is a security vulnerability.
+
+    You can change this behaviour by:
+
+    .. code-block:: python
+
+       from urllib3 import Retry
+       Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT = frozenset()
+       Retry.DEFAULT.remove_headers_on_redirect = Retry.DEFAULT_REMOVE_HEADERS_ON_REDIRECT
+
+.. marker-proxy-end
+
 Delete data
 ^^^^^^^^^^^
 .. marker-delete-start

docs/usage.rst

@@ -10,18 +10,18 @@ Query
   :start-after: marker-query-start
   :end-before: marker-query-end
 
-Pandas DataFrame
-^^^^^^^^^^^^^^^^
-.. include:: ../README.rst
-  :start-after: marker-pandas-start
-  :end-before: marker-pandas-end
-
 Write
 ^^^^^
 .. include:: ../README.rst
   :start-after: marker-writes-start
   :end-before: marker-writes-end
 
+Pandas DataFrame
+^^^^^^^^^^^^^^^^
+.. include:: ../README.rst
+  :start-after: marker-pandas-start
+  :end-before: marker-pandas-end
+
 Delete data
 ^^^^^^^^^^^
 .. include:: ../README.rst
@@ -34,6 +34,12 @@ Gzip support
   :start-after: marker-gzip-start
   :end-before: marker-gzip-end
 
+Proxy configuration
+^^^^^^^^^^^^^^^^^^^
+.. include:: ../README.rst
+  :start-after: marker-proxy-start
+  :end-before: marker-proxy-end
+
 Debugging
 ^^^^^^^^^
 

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