feat: support gzip_threshold

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -21,6 +21,7 @@
 INFLUX_ORG = "INFLUX_ORG"
 INFLUX_PRECISION = "INFLUX_PRECISION"
 INFLUX_AUTH_SCHEME = "INFLUX_AUTH_SCHEME"
+INFLUX_GZIP_THRESHOLD = "INFLUX_GZIP_THRESHOLD"
 
 
 def write_client_options(**kwargs):
@@ -114,6 +115,30 @@ def _parse_precision(precision):
     return precision
 
 
+def _parse_gzip_threshold(threshold):
+    """
+    Parses and validates the provided threshold value.
+
+    This function ensures that the given threshold is a valid integer value,
+    and it raises an appropriate error if the threshold is not valid. It also
+    enforces that the threshold value is non-negative.
+
+    :param threshold: The input threshold value to be parsed and validated.
+    :type threshold: Any
+    :return: The validated threshold value as an integer.
+    :rtype: int
+    :raises ValueError: If the provided threshold is not an integer or if it is
+        negative.
+    """
+    try:
+        threshold = int(threshold)
+    except (TypeError, ValueError):
+        raise ValueError(f"Invalid threshold value: {threshold}. Must be integer.")
+    if threshold < 0:
+        raise ValueError(f"Invalid threshold value: {threshold}. Must be non-negative.")
+    return threshold
+
+
 class InfluxDBClient3:
     def __init__(
             self,
@@ -254,6 +279,11 @@ def from_env(cls, **kwargs: Any) -> 'InfluxDBClient3':
 
         write_options = WriteOptions(write_type=WriteType.synchronous)
 
+        gzip_threshold = os.getenv(INFLUX_GZIP_THRESHOLD)
+        if gzip_threshold is not None:
+            kwargs['gzip_threshold'] = _parse_gzip_threshold(gzip_threshold)
+            kwargs['enable_gzip'] = True
+
         precision = os.getenv(INFLUX_PRECISION)
         if precision is not None:
             write_options.write_precision = _parse_precision(precision)

influxdb_client_3/write_client/_sync/api_client.py

@@ -92,6 +92,25 @@ def set_default_header(self, header_name, header_value):
         """Set HTTP header for this API client."""
         self.default_headers[header_name] = header_value
 
+    @staticmethod
+    def should_gzip(payload: str, enable_gzip: bool = True, gzip_threshold: int = None) -> bool:
+        """
+        Determine if the payload should be compressed with gzip.
+
+        Args:
+            payload: The string content to potentially compress
+            enable_gzip: Flag indicating if gzip compression is enabled
+            gzip_threshold: Minimum size in bytes for compression to be applied
+
+        Returns:
+            bool: True if the payload should be compressed, False otherwise
+        """
+        if not enable_gzip:
+            return False
+
+        payload_size = len(payload.encode('utf-8'))
+        return gzip_threshold is not None and payload_size >= gzip_threshold
+
     def __call_api(
             self, resource_path, method, path_params=None,
             query_params=None, header_params=None, body=None, post_params=None,
@@ -102,9 +121,16 @@ def __call_api(
         config = self.configuration
         self._signin(resource_path=resource_path)
 
+        # body
+        should_gzip = False
+        if body:
+            should_gzip = self.should_gzip(config.enable_gzip, config.gzip_threshold, body)
+            body = self.sanitize_for_serialization(body)
+            body = config.update_request_body(resource_path, body, should_gzip)
+
         # header parameters
         header_params = header_params or {}
-        config.update_request_header_params(resource_path, header_params)
+        config.update_request_header_params(resource_path, header_params, should_gzip)
         header_params.update(self.default_headers)
         if self.cookie:
             header_params['Cookie'] = self.cookie
@@ -141,11 +167,6 @@ def __call_api(
         # auth setting
         self.update_params_for_auth(header_params, query_params, auth_settings)
 
-        # body
-        if body:
-            body = self.sanitize_for_serialization(body)
-            body = config.update_request_body(resource_path, body)
-
         # request url
         url = self.configuration.host + resource_path
 

influxdb_client_3/write_client/client/_base.py

@@ -49,6 +49,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 = kwargs.get('gzip_threshold', None)
         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)
@@ -277,9 +278,9 @@ def __init__(self):
         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:
+    def update_request_header_params(self, path: str, params: dict, should_gzip: bool = False):
+        super().update_request_header_params(path, params, should_gzip)
+        if should_gzip:
             # GZIP Request
             if path == '/api/v2/write':
                 params["Content-Encoding"] = "gzip"
@@ -293,9 +294,9 @@ def update_request_header_params(self, path: str, params: dict):
             pass
         pass
 
-    def update_request_body(self, path: str, body):
-        _body = super().update_request_body(path, body)
-        if self.enable_gzip:
+    def update_request_body(self, path: str, body, should_gzip: bool = False):
+        _body = super().update_request_body(path, body, should_gzip)
+        if should_gzip:
             # GZIP Request
             if path == '/api/v2/write':
                 import gzip

influxdb_client_3/write_client/client/influxdb_client.py

@@ -16,7 +16,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,
-                 default_tags: dict = None, **kwargs) -> None:
+                 default_tags: dict = None, gzip_threshold=None, **kwargs) -> None:
         """
         Initialize defaults.
 
@@ -52,7 +52,8 @@ 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/configuration.py

@@ -98,6 +98,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.
@@ -245,19 +249,21 @@ def to_debug_report(self):
                "SDK Package Version: {client_version}".\
                format(env=sys.platform, pyversion=sys.version, client_version=VERSION)
 
-    def update_request_header_params(self, path: str, params: dict):
+    def update_request_header_params(self, path: str, params: dict, should_gzip: bool = False):
         """Update header params based on custom settings.
 
-        :param path: Resource path
+        :param path: Resource path.
         :param params: Header parameters dict to be updated.
+        :param should_gzip: Describes if request body should be gzip compressed.
         """
         pass
 
-    def update_request_body(self, path: str, body):
+    def update_request_body(self, path: str, body, should_gzip: bool = False):
         """Update http body based on custom settings.
 
-        :param path: Resource path
+        :param path: Resource path.
         :param body: Request body to be updated.
+        :param should_gzip: Describes if request body should be gzip compressed.
         :return: Updated body
         """
         return body

tests/test_api_client.py

@@ -139,3 +139,21 @@ 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_should_gzip(self):
+        # Test when gzip is disabled
+        assert not ApiClient.should_gzip("test", enable_gzip=False, gzip_threshold=1)
+
+        # Test when threshold is None
+        assert not ApiClient.should_gzip("test", enable_gzip=True, gzip_threshold=None)
+
+        # Test payload smaller than threshold
+        assert not ApiClient.should_gzip("test", enable_gzip=True, gzip_threshold=10000)
+
+        # Test payload larger than threshold
+        large_payload = "x" * 10000
+        assert ApiClient.should_gzip(large_payload, enable_gzip=True, gzip_threshold=1000)
+
+        # Test exact threshold match
+        exact_payload = "x" * 1000
+        assert ApiClient.should_gzip(exact_payload, enable_gzip=True, gzip_threshold=1000)