refactor: write rest client

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -1,3 +1,5 @@
+import multiprocessing
+
 import importlib.util
 import json
 import os
@@ -326,7 +328,8 @@ def __init__(
             ssl_context=kwargs.get('ssl_context', None),
             proxy=kwargs.get('proxy', None),
             proxy_headers=kwargs.get('proxy_headers', None),
-            retries=kwargs.get('retries', False)
+            retries=kwargs.get('retries', False),
+            connection_pool_maxsize=kwargs.get('connection_pool_maxsize', multiprocessing.cpu_count() * 5,)
         )
 
         if point_settings is None:
@@ -721,9 +724,13 @@ def get_server_version(self) -> Optional[str]:
                 return value
 
         string_body = resp.get_string_body()
-        if string_body is not None:
-            return json.loads(string_body)['version']
-        return None
+        if not string_body:
+            return None
+        try:
+            data = json.loads(string_body)
+        except (ValueError, TypeError):
+            return None
+        return data.get('version')
 
     def flush(self):
         """
@@ -744,6 +751,8 @@ def close(self):
             self._write_api.close()
         if self._query_api is not None:
             self._query_api.close()
+        if self.rest_client is not None:
+            self.rest_client.close()
 
     def __enter__(self):
         return self

influxdb_client_3/write_client/_sync/rest_client.py

@@ -155,15 +155,19 @@ def request(self, method, url, query_params=None, headers=None,
                                       **urlopen_kw)
 
         r = RESTResponse(r)
-        r.data = r.data.decode('utf8')
+        if isinstance(r.data, (bytes, bytearray)):
+            r.data = r.data.decode('utf-8')
+        elif r.data is None:
+            r.data = ''
 
         if not 200 <= r.status <= 299:
-            print('hahaha')
-            print(r.data)
             raise ApiException(http_resp=r)
 
         return r
 
+    def close(self):
+        self.pool_manager.close()
+
     def __getstate__(self):
         """Return a dict of attributes that you want to pickle."""
         state = self.__dict__.copy()

influxdb_client_3/write_client/client/write_api.py

@@ -319,7 +319,7 @@ def __init__(self,
         :param enable_gzip: Boolean indicating whether GZIP compression is enabled.
         :param auth_scheme: The authentication scheme to use for the connection
             (e.g., Basic, Bearer).
-        :param timeout: Timeout duration (in seconds) for HTTP requests.
+        :param timeout: Timeout duration (in milliseconds) for HTTP requests.
         :param pool_threads: Number of threads used for connection pools.
         :param default_header: Default HTTP headers to include in every request.
         :param rest_client: An instance of a RestClient for internal HTTP communication.
@@ -536,6 +536,10 @@ def flush(self):
 
     def close(self):
         """Flush data and dispose a batching buffer."""
+        if self._pool is not None:
+            self._pool.close()
+            self._pool.join()
+            self._pool = None
         if self._subject is None:
             return  # Already closed
 
@@ -1088,14 +1092,6 @@ def _on_next(self, response: _BatchResponse):
                 except Exception as e:
                     logger.error("The configured success callback threw an exception: %s", e)
 
-    def __getstate__(self):
-        """Return a dict of attributes that you want to pickle."""
-        state = self.__dict__.copy()
-        # Remove rx
-        del state['_subject']
-        del state['_disposable']
-        return state
-
     def __setstate__(self, state):
         """Set your object with the provided dict."""
         self.__dict__.update(state)

tests/test_write_api.py

@@ -249,7 +249,7 @@ def test_api_error_headers(self):
             )
         )
         with self.assertRaises(InfluxDBError) as err:
-            client._write_api.write("TEST_ORG", "TEST_BUCKET", "data,foo=bar val=3.14")
+            client._write_api.write("TEST_BUCKET", "TEST_ORG", "data,foo=bar val=3.14")
         self.assertEqual(body_dic['error'], err.exception.message)
         headers = err.exception.getheaders()
         self.assertEqual(4, len(headers))
@@ -290,7 +290,7 @@ def test_request_arg_timeout(self, mock_request):
         )
 
         with pytest.raises(ConnectTimeoutError):
-            client._write_api.write("TEST_ORG", "TEST_BUCKET", "data,foo=bar val=3.14",
+            client._write_api.write("TEST_BUCKET", "TEST_ORG", "data,foo=bar val=3.14",
                                     _request_timeout=100)
         self.assertEqual(0.1, self.received_timeout_total)
         self.received_timeout_total = None