feat: allow to configure a connection pool maxsize (influxdata#215)

bednar · web-flow · commit a4e971d2d425 · 2021-03-26T11:55:56.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,8 @@
 ## 1.16.0 [unreleased]
 
 ### Features
-1. [#203](https://github.com/influxdata/influxdb-client-python/pull/203): Allow configuring client via TOML file.
+1. [#203](https://github.com/influxdata/influxdb-client-python/pull/203): Configure a client via TOML file
+1. [#215](https://github.com/influxdata/influxdb-client-python/pull/215): Configure a connection pool maxsize
 
 ### Bug Fixes
 1. [#206](https://github.com/influxdata/influxdb-client-python/pull/207): Use default (system) certificates instead of Mozilla's root certificates (certifi.where())
diff --git a/README.rst b/README.rst
@@ -174,6 +174,7 @@ The following options are supported:
 - ``timeout`` - socket timeout in ms (default value is 10000)
 - ``verify_ssl`` - set this to false to skip verifying SSL certificate when calling API from https server
 - ``ssl_ca_cert`` - set this to customize the certificate file to verify the peer
+- ``connection_pool_maxsize`` - set the number of connections to save that can be reused by urllib3
 
 .. code-block:: python
 
@@ -200,6 +201,7 @@ Supported properties are:
 - ``INFLUXDB_V2_TIMEOUT`` - socket timeout in ms (default value is 10000)
 - ``INFLUXDB_V2_VERIFY_SSL`` - set this to false to skip verifying SSL certificate when calling API from https server
 - ``INFLUXDB_V2_SSL_CA_CERT`` - set this to customize the certificate file to verify the peer
+- ``INFLUXDB_V2_CONNECTION_POOL_MAXSIZE`` - set this to customize the certificate file to verify the peer
 
 .. code-block:: python
 
diff --git a/influxdb_client/client/influxdb_client.py b/influxdb_client/client/influxdb_client.py
@@ -35,6 +35,8 @@ def __init__(self, url, token, debug=None, timeout=10000, enable_gzip=False, org
         :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 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
                                                except batching writes. As a default there is no one retry strategy.
 
@@ -56,6 +58,7 @@ def __init__(self, url, token, debug=None, timeout=10000, enable_gzip=False, org
         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.connection_pool_maxsize = kwargs.get('connection_pool_maxsize', conf.connection_pool_maxsize)
 
         auth_token = self.token
         auth_header_name = "Authorization"
@@ -82,6 +85,7 @@ def from_config_file(cls, config_file: str = "config.ini", debug=None, enable_gz
             - timeout,
             - verify_ssl
             - ssl_ca_cert
+            - connection_pool_maxsize
 
         config.ini example::
 
@@ -90,6 +94,7 @@ def from_config_file(cls, config_file: str = "config.ini", debug=None, enable_gz
             org=my-org
             token=my-token
             timeout=6000
+            connection_pool_maxsize=25
 
             [tags]
             id = 132-987-655
@@ -103,6 +108,7 @@ def from_config_file(cls, config_file: str = "config.ini", debug=None, enable_gz
                 token = "my-token"
                 org = "my-org"
                 timeout = 6000
+                connection_pool_maxsize = 25
 
             [tags]
                 id = "132-987-655"
@@ -137,18 +143,19 @@ def config_value(key: str):
         if config.has_option('influx2', 'ssl_ca_cert'):
             ssl_ca_cert = config_value('ssl_ca_cert')
 
+        connection_pool_maxsize = None
+        if config.has_option('influx2', 'connection_pool_maxsize'):
+            connection_pool_maxsize = config_value('connection_pool_maxsize')
+
         default_tags = None
 
         if config.has_section('tags'):
             tags = {k: v.strip('"') for k, v in config.items('tags')}
             default_tags = dict(tags)
 
-        if timeout:
-            return cls(url, token, debug=debug, timeout=int(timeout), org=org, default_tags=default_tags,
-                       enable_gzip=enable_gzip, verify_ssl=_to_bool(verify_ssl), ssl_ca_cert=ssl_ca_cert)
-
-        return cls(url, token, debug=debug, org=org, default_tags=default_tags, enable_gzip=enable_gzip,
-                   verify_ssl=_to_bool(verify_ssl), ssl_ca_cert=ssl_ca_cert)
+        return cls(url, token, debug=debug, timeout=_to_int(timeout), org=org, default_tags=default_tags,
+                   enable_gzip=enable_gzip, verify_ssl=_to_bool(verify_ssl), ssl_ca_cert=ssl_ca_cert,
+                   connection_pool_maxsize=_to_int(connection_pool_maxsize))
 
     @classmethod
     def from_env_properties(cls, debug=None, enable_gzip=False):
@@ -162,22 +169,25 @@ def from_env_properties(cls, debug=None, enable_gzip=False):
             - INFLUXDB_V2_TIMEOUT
             - INFLUXDB_V2_VERIFY_SSL
             - INFLUXDB_V2_SSL_CA_CERT
+            - INFLUXDB_V2_CONNECTION_POOL_MAXSIZE
         """
         url = os.getenv('INFLUXDB_V2_URL', "http://localhost:8086")
         token = os.getenv('INFLUXDB_V2_TOKEN', "my-token")
         timeout = os.getenv('INFLUXDB_V2_TIMEOUT', "10000")
         org = os.getenv('INFLUXDB_V2_ORG', "my-org")
         verify_ssl = os.getenv('INFLUXDB_V2_VERIFY_SSL', "True")
         ssl_ca_cert = os.getenv('INFLUXDB_V2_SSL_CA_CERT', None)
+        connection_pool_maxsize = os.getenv('INFLUXDB_V2_CONNECTION_POOL_MAXSIZE', None)
 
         default_tags = dict()
 
         for key, value in os.environ.items():
             if key.startswith("INFLUXDB_V2_TAG_"):
                 default_tags[key[16:].lower()] = value
 
-        return cls(url, token, debug=debug, timeout=int(timeout), org=org, default_tags=default_tags,
-                   enable_gzip=enable_gzip, verify_ssl=_to_bool(verify_ssl), ssl_ca_cert=ssl_ca_cert)
+        return cls(url, token, debug=debug, timeout=_to_int(timeout), org=org, default_tags=default_tags,
+                   enable_gzip=enable_gzip, verify_ssl=_to_bool(verify_ssl), ssl_ca_cert=ssl_ca_cert,
+                   connection_pool_maxsize=_to_int(connection_pool_maxsize))
 
     def write_api(self, write_options=WriteOptions(), point_settings=PointSettings()) -> WriteApi:
         """
@@ -322,5 +332,9 @@ def update_request_body(self, path: str, body):
         return _body
 
 
-def _to_bool(verify_ssl):
-    return str(verify_ssl).lower() in ("yes", "true")
+def _to_bool(bool_value):
+    return str(bool_value).lower() in ("yes", "true")
+
+
+def _to_int(int_value):
+    return int(int_value) if int_value is not None else None
diff --git a/tests/config.ini b/tests/config.ini
@@ -3,6 +3,7 @@ url=http://localhost:8086
 org=my-org
 token=my-token
 timeout=6000
+connection_pool_maxsize=55
 
 [tags]
 id = 132-987-655
diff --git a/tests/config.toml b/tests/config.toml
@@ -4,6 +4,7 @@
   org = "my-org"
   active = true
   timeout = 6000
+  connection_pool_maxsize = 55
 
 [tags]
   id = "132-987-655"
diff --git a/tests/test_InfluxDBClient.py b/tests/test_InfluxDBClient.py
@@ -17,6 +17,10 @@ def tearDown(self) -> None:
         if hasattr(self, 'httpd_thread'):
             self.httpd_thread.join()
 
+    def test_default_conf(self):
+        self.client = InfluxDBClient(url="http://localhost:8086", token="my-token", org="my-org")
+        self.assertIsNotNone(self.client.api_client.configuration.connection_pool_maxsize)
+
     def test_TrailingSlashInUrl(self):
         self.client = InfluxDBClient(url="http://localhost:8086", token="my-token", org="my-org")
         self.assertEqual('http://localhost:8086', self.client.api_client.configuration.host)
@@ -66,6 +70,7 @@ def assertConfig(self):
         self.assertEqual("132-987-655", self.client.default_tags["id"])
         self.assertEqual("California Miner", self.client.default_tags["customer"])
         self.assertEqual("${env.data_center}", self.client.default_tags["data_center"])
+        self.assertEqual(55, self.client.api_client.configuration.connection_pool_maxsize)
 
     def test_init_from_file_ssl_default(self):
         self.client = InfluxDBClient.from_config_file(f'{os.path.dirname(__file__)}/config.ini')
@@ -113,6 +118,12 @@ def test_init_from_env_ssl_ca_cert(self):
 
         self.assertEqual("/my/custom/path/to/cert", self.client.api_client.configuration.ssl_ca_cert)
 
+    def test_init_from_env_connection_pool_maxsize(self):
+        os.environ["INFLUXDB_V2_CONNECTION_POOL_MAXSIZE"] = "29"
+        self.client = InfluxDBClient.from_env_properties()
+
+        self.assertEqual(29, self.client.api_client.configuration.connection_pool_maxsize)
+
     def _start_http_server(self):
         import http.server
         import ssl
