feat: support env

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -96,7 +96,9 @@ def from_env(**kwargs: Any) -> 'InfluxDBClient3':
     write_options = WriteOptions(write_type=WriteType.synchronous)
 
     if os.getenv(INFLUX_GZIP_THRESHOLD) is not None:
-        write_options.gzip_threshold = int(os.getenv(INFLUX_GZIP_THRESHOLD))
+        gzip_threshold = int(os.getenv(INFLUX_GZIP_THRESHOLD))
+        write_options.enable_gzip = True
+        write_options.gzip_threshold = gzip_threshold
 
     if os.getenv(INFLUX_PRECISION) is not None:
         write_options.write_precision = os.getenv(INFLUX_PRECISION)
@@ -105,7 +107,6 @@ def from_env(**kwargs: Any) -> 'InfluxDBClient3':
 
     if os.getenv(INFLUX_AUTH_SCHEME) is not None:
         kwargs['auth_scheme'] = os.getenv(INFLUX_AUTH_SCHEME)
-
     org = os.getenv(INFLUX_ORG, "default")
 
     return InfluxDBClient3(
@@ -210,17 +211,18 @@ def __init__(
 
         write_type = DefaultWriteOptions['write_type']
         write_precision = DefaultWriteOptions['write_precision']
-        gzip_threshold = DefaultWriteOptions['gzip_threshold']
+        gzip_threshold = None
         if isinstance(write_client_options, dict) and write_client_options.get('write_options') is not None:
             write_opts = write_client_options['write_options']
-            write_type = getattr(write_opts, 'write_type')
-            write_precision = getattr(write_opts, 'write_precision')
+            write_type = getattr(write_opts, 'write_type', write_type)
+            write_precision = getattr(write_opts, 'write_precision', write_precision)
             gzip_threshold = getattr(write_opts, 'gzip_threshold')
 
         write_options = WriteOptions(
             write_type=write_type,
             write_precision=write_precision,
             gzip_threshold=gzip_threshold,
+            enable_gzip=kwargs.get('enable_gzip', False)
         )
 
         self._write_client_options = {
@@ -244,6 +246,8 @@ def __init__(
             url=f"{scheme}://{hostname}:{port}",
             token=self._token,
             org=self._org,
+            enable_gzip=write_options.enable_gzip,
+            gzip_threshold=write_options.gzip_threshold,
             **kwargs)
 
         self._write_api = _WriteApi(influxdb_client=self._client, **self._write_client_options)

influxdb_client_3/write_client/_sync/api_client.py

@@ -120,6 +120,10 @@ def __call_api(
         config = self.configuration
         self._signin(resource_path=resource_path)
 
+        gzip_threshold = config.gzip_threshold
+        enable_gzip = config.enable_gzip
+        self.should_compress = self.check_should_compress(body, gzip_threshold, enable_gzip)
+        
         # header parameters
         header_params = header_params or {}
         config.update_request_header_params(resource_path, header_params)
@@ -192,6 +196,12 @@ def __call_api(
             return (return_data, response_data.status,
                     response_data.getheaders())
 
+    def check_should_compress(self, body: bytearray, gzip_threshold: int, enable_gzip: bool) -> bool:
+        body_size = len(body)
+        if enable_gzip is True or (enable_gzip is not False and (gzip_threshold and body_size >= gzip_threshold)):
+            return True
+        return False
+
     def sanitize_for_serialization(self, obj):
         """Build a JSON POST object.
 

influxdb_client_3/write_client/client/_base.py

@@ -34,7 +34,7 @@
 
 # noinspection PyMethodMayBeStatic
 class _BaseClient(object):
-    def __init__(self, url, token, debug=None, timeout=10_000, enable_gzip=False, org: str = None,
+    def __init__(self, url, token, debug=None, timeout=10_000, enable_gzip=False, gzip_threshold=None, org: str = None,
                  default_tags: dict = None, http_client_logger: str = None, **kwargs) -> None:
         self.url = url
         self.org = org
@@ -47,6 +47,7 @@ def __init__(self, url, token, debug=None, timeout=10_000, enable_gzip=False, or
         else:
             self.conf.host = self.url
         self.conf.enable_gzip = enable_gzip
+        self.conf.gzip_threshold = gzip_threshold
         self.conf.verify_ssl = kwargs.get('verify_ssl', True)
         self.conf.ssl_ca_cert = kwargs.get('ssl_ca_cert', None)
         self.conf.cert_file = kwargs.get('cert_file', None)
@@ -271,12 +272,14 @@ class _Configuration(Configuration):
     def __init__(self):
         Configuration.__init__(self)
         self.enable_gzip = False
+        self.gzip_threshold = None
+        self.should_compress = False
         self.username = None
         self.password = None
 
     def update_request_header_params(self, path: str, params: dict):
         super().update_request_header_params(path, params)
-        if self.enable_gzip:
+        if self.should_compress:
             # GZIP Request
             if path == '/api/v2/write':
                 params["Content-Encoding"] = "gzip"
@@ -292,7 +295,7 @@ def update_request_header_params(self, path: str, params: dict):
 
     def update_request_body(self, path: str, body):
         _body = super().update_request_body(path, body)
-        if self.enable_gzip:
+        if self.should_compress:
             # GZIP Request
             if path == '/api/v2/write':
                 import gzip

influxdb_client_3/write_client/client/influxdb_client.py

@@ -13,7 +13,7 @@
 class InfluxDBClient(_BaseClient):
     """InfluxDBClient is client for InfluxDB v2."""
 
-    def __init__(self, url, token: str = None, debug=None, timeout=10_000, enable_gzip=False, org: str = None,
+    def __init__(self, url, token: str = None, debug=None, timeout=10_000, enable_gzip=False, gzip_threshold=None, org: str = None,
                  default_tags: dict = None, **kwargs) -> None:
         """
         Initialize defaults.
@@ -50,7 +50,7 @@ def __init__(self, url, token: str = None, debug=None, timeout=10_000, enable_gz
         :key str password: ``password`` to authenticate via username and password credentials to the InfluxDB 2.x
         :key list[str] profilers: list of enabled Flux profilers
         """
-        super().__init__(url=url, token=token, debug=debug, timeout=timeout, enable_gzip=enable_gzip, org=org,
+        super().__init__(url=url, token=token, debug=debug, timeout=timeout, enable_gzip=enable_gzip, gzip_threshold=gzip_threshold, org=org,
                          default_tags=default_tags, http_client_logger="urllib3", **kwargs)
 
         from influxdb_client_3.write_client._sync.api_client import ApiClient

influxdb_client_3/write_client/client/write_api.py

@@ -38,11 +38,12 @@ class WriteType(Enum):
     asynchronous = 2
     synchronous = 3
 
+DEFAULT_GZIP_THRESHOLD = 1000
 
+#todo: convert to enum
 DefaultWriteOptions = {
     'write_type': WriteType.synchronous,
-    'write_precision': WritePrecision.NS,
-    'gzip_threshold': 1000
+    'write_precision': WritePrecision.NS
 }
 
 
@@ -59,7 +60,8 @@ def __init__(self, write_type: WriteType = WriteType.batching,
                  exponential_base=2,
                  max_close_wait=300_000,
                  write_precision=DEFAULT_WRITE_PRECISION,
-                 gzip_threshold=1000,
+                 gzip_threshold=None,
+                 enable_gzip=False,
                  write_scheduler=ThreadPoolScheduler(max_workers=1)) -> None:
         """
         Create write api configuration.
@@ -92,6 +94,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
         self.max_close_wait = max_close_wait
         self.write_precision = write_precision
         self.gzip_threshold = gzip_threshold
+        self.enable_gzip = enable_gzip
 
     def to_retry_strategy(self, **kwargs):
         """

influxdb_client_3/write_client/configuration.py

@@ -49,6 +49,7 @@ class Configuration(object, metaclass=TypeWithDefault):
     Ref: https://openapi-generator.tech
     Do not edit the class manually.
     """
+    #todo: remove wrong document
 
     def __init__(self):
         """Initialize configuration."""
@@ -118,6 +119,10 @@ def __init__(self):
         # Safe chars for path_param
         self.safe_chars_for_path_param = ''
 
+        # Compression settings
+        self.enable_gzip = False
+        self.gzip_threshold = None
+
     @property
     def logger_file(self):
         """Logger file.

tests/test_api_client.py

@@ -139,3 +139,48 @@ def test_api_error_headers(self):
         self.assertEqual(headers['Trace-Sampled'], 'false')
         self.assertEqual(headers['X-Influxdb-Request-Id'], requestid)
         self.assertEqual(headers['X-Influxdb-Build'], 'Mock')
+
+    def test_check_should_compress_true(self):
+        conf = Configuration()
+        client = ApiClient(conf)
+
+        body = bytearray("12345678901234567890".encode("utf-8")) # len = 20
+        tests = [
+            {
+                'gzip_threshold': 10,
+                'enable_gzip': True,
+                'expected': True
+            },
+            {
+                'gzip_threshold': 30,
+                'enable_gzip': True,
+                'expected': True
+            },
+            {
+                'gzip_threshold': None,
+                'enable_gzip': True,
+                'expected': True
+            },
+            {
+                'gzip_threshold': 30,
+                'enable_gzip': None,
+                'expected': False
+            },
+            {
+                'gzip_threshold': 30,
+                'enable_gzip': False,
+                'expected': False
+            },
+            {
+                'gzip_threshold': 10,
+                'enable_gzip': None,
+                'expected': True
+            },
+        ]
+
+        for test in tests:
+            gzip_threshold = test['gzip_threshold']
+            enable_gzip = test['enable_gzip']
+            expected = test['expected']
+            result = client.check_should_compress(body, gzip_threshold, enable_gzip)
+            self.assertEqual(result, expected)

tests/test_influxdb_client_3.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import patch
 
-from influxdb_client_3 import InfluxDBClient3, from_env, WritePrecision, WriteType
+from influxdb_client_3 import InfluxDBClient3, from_env, WritePrecision, DefaultWriteOptions
 from tests.util import asyncio_run
 from tests.util.mocks import ConstantFlightServer, ConstantData
 
@@ -75,13 +75,17 @@ async def test_query_async(self):
             assert {'data': 'query_type', 'reference': 'sql', 'value': -1.0} in result_list
 
     def test_default_client(self):
-        expected_precision = WritePrecision.NS
-        expected_write_type = WriteType.synchronous
+        expected_precision = DefaultWriteOptions['write_precision']
+        expected_write_type = DefaultWriteOptions['write_type']
+        expected_gzip_threshold = None
+        expected_gzip_enabled = False
 
         def verify_client_write_options(c):
             write_options = c._write_client_options.get('write_options')
             self.assertEqual(write_options.write_precision, expected_precision)
             self.assertEqual(write_options.write_type, expected_write_type)
+            self.assertEqual(write_options.gzip_threshold, expected_gzip_threshold)
+            self.assertEqual(write_options.enable_gzip, expected_gzip_enabled)
 
             self.assertEqual(c._write_api._write_options.write_precision, expected_precision)
             self.assertEqual(c._write_api._write_options.write_type, expected_write_type)