feat: disable grpc response compression

Author: jansimonb

influxdb_client_3/__init__.py

@@ -31,6 +31,7 @@
 INFLUX_WRITE_NO_SYNC = "INFLUX_WRITE_NO_SYNC"
 INFLUX_WRITE_TIMEOUT = "INFLUX_WRITE_TIMEOUT"
 INFLUX_QUERY_TIMEOUT = "INFLUX_QUERY_TIMEOUT"
+INFLUX_DISABLE_GRPC_COMPRESSION = "INFLUX_DISABLE_GRPC_COMPRESSION"
 
 
 def write_client_options(**kwargs):
@@ -190,6 +191,7 @@ def __init__(
             flight_client_options=None,
             write_port_overwrite=None,
             query_port_overwrite=None,
+            disable_grpc_compression=False,
             **kwargs):
         """
         Initialize an InfluxDB client.
@@ -206,6 +208,8 @@ def __init__(
         :type write_client_options: dict[str, any]
         :param flight_client_options: dictionary for providing additional arguments for the FlightClient.
         :type flight_client_options: dict[str, any]
+        :param disable_grpc_compression: Disable gRPC compression for Flight query responses. Default is False.
+        :type disable_grpc_compression: bool
         :key auth_scheme: token authentication scheme. Set to "Bearer" for Edge.
         :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.
@@ -291,6 +295,8 @@ def __init__(
             connection_string = f"grpc+tcp://{hostname}:{port}"
 
         q_opts_builder = QueryApiOptionsBuilder()
+        if disable_grpc_compression:
+            q_opts_builder.disable_grpc_compression(True)
         if kw_keys.__contains__('ssl_ca_cert'):
             q_opts_builder.root_certs(kwargs.get('ssl_ca_cert', None))
         if kw_keys.__contains__('verify_ssl'):
@@ -361,13 +367,20 @@ def from_env(cls, **kwargs: Any) -> 'InfluxDBClient3':
         if os.getenv(INFLUX_AUTH_SCHEME) is not None:
             kwargs['auth_scheme'] = os.getenv(INFLUX_AUTH_SCHEME)
 
+        disable_grpc_compression = os.getenv(INFLUX_DISABLE_GRPC_COMPRESSION)
+        if disable_grpc_compression is not None:
+            disable_grpc_compression = disable_grpc_compression.strip().lower() in ['true', '1', 't', 'y', 'yes']
+        else:
+            disable_grpc_compression = False
+
         org = os.getenv(INFLUX_ORG, "default")
         return InfluxDBClient3(
             host=required_vars[INFLUX_HOST],
             token=required_vars[INFLUX_TOKEN],
             database=required_vars[INFLUX_DATABASE],
             write_client_options=write_client_option,
             org=org,
+            disable_grpc_compression=disable_grpc_compression,
             **kwargs
         )
 

influxdb_client_3/query/query_api.py

@@ -19,19 +19,22 @@ class QueryApiOptions(object):
     proxy (str): URL to a proxy server
     flight_client_options (dict): base set of flight client options passed to internal pyarrow.flight.FlightClient
     timeout(float): timeout in seconds to wait for a response
+    disable_grpc_compression (bool): disable gRPC compression for query responses
     """
     _DEFAULT_TIMEOUT = 300.0
     tls_root_certs: bytes = None
     tls_verify: bool = None
     proxy: str = None
     flight_client_options: dict = None
     timeout: float = None
+    disable_grpc_compression: bool = False
 
     def __init__(self, root_certs_path: str,
                  verify: bool,
                  proxy: str,
                  flight_client_options: dict,
-                 timeout: float = _DEFAULT_TIMEOUT):
+                 timeout: float = _DEFAULT_TIMEOUT,
+                 disable_grpc_compression: bool = False):
         """
         Initialize a set of QueryApiOptions
 
@@ -41,13 +44,15 @@ def __init__(self, root_certs_path: str,
         :param flight_client_options: set of flight_client_options
                to be passed to internal pyarrow.flight.FlightClient.
         :param timeout: timeout in seconds to wait for a response.
+        :param disable_grpc_compression: disable gRPC compression for query responses.
         """
         if root_certs_path:
             self.tls_root_certs = self._read_certs(root_certs_path)
         self.tls_verify = verify
         self.proxy = proxy
         self.flight_client_options = flight_client_options
         self.timeout = timeout
+        self.disable_grpc_compression = disable_grpc_compression
 
     def _read_certs(self, path: str) -> bytes:
         with open(path, "rb") as certs_file:
@@ -75,6 +80,7 @@ class QueryApiOptionsBuilder(object):
     _proxy: str = None
     _flight_client_options: dict = None
     _timeout: float = None
+    _disable_grpc_compression: bool = False
 
     def root_certs(self, path: str):
         self._root_certs_path = path
@@ -96,6 +102,11 @@ def timeout(self, timeout: float):
         self._timeout = timeout
         return self
 
+    def disable_grpc_compression(self, disable: bool):
+        """Disable gRPC compression for query responses."""
+        self._disable_grpc_compression = disable
+        return self
+
     def build(self) -> QueryApiOptions:
         """Build a QueryApiOptions object with previously set values"""
         return QueryApiOptions(
@@ -104,6 +115,7 @@ def build(self) -> QueryApiOptions:
             proxy=self._proxy,
             flight_client_options=self._flight_client_options,
             timeout=self._timeout,
+            disable_grpc_compression=self._disable_grpc_compression,
         )
 
 
@@ -162,6 +174,13 @@ def __init__(self,
                 self._flight_client_options["disable_server_verification"] = not options.tls_verify
             if options.timeout is not None:
                 self._default_timeout = options.timeout
+            if options.disable_grpc_compression:
+                # Disable gRPC response compression by only enabling identity algorithm
+                # Bitset: bit 0 = identity, bit 1 = deflate, bit 2 = gzip
+                # Setting to 1 (0b001) enables only identity (no compression)
+                self._flight_client_options["generic_options"].append(
+                    ("grpc.compression_enabled_algorithms_bitset", 1)
+                )
         if self._proxy:
             self._flight_client_options["generic_options"].append(("grpc.http_proxy", self._proxy))
         self._flight_client = FlightClient(connection_string, **self._flight_client_options)

setup.py

@@ -51,7 +51,12 @@ def get_version():
         'pandas': ['pandas'],
         'polars': ['polars'],
         'dataframe': ['pandas', 'polars'],
-        'test': ['pytest', 'pytest-cov', 'pytest-httpserver']
+        'test': [
+            'pytest',
+            'pytest-cov',
+            'pytest-httpserver',
+            'mitmproxy>=11.0.0,<12.0.0; python_version>="3.10"',
+        ]
     },
     install_requires=requires,
     python_requires='>=3.8',

tests/test_influxdb_client_3.py

@@ -302,6 +302,74 @@ def test_parse_invalid_write_timeout_range(self):
         with self.assertRaisesRegex(ValueError, ".*Must be non-negative.*"):
             InfluxDBClient3.from_env()
 
+    def assertGrpcCompressionDisabled(self, client, disabled):
+        """Assert whether gRPC compression is disabled for the client."""
+        self.assertIsInstance(client, InfluxDBClient3)
+        generic_options = dict(client._query_api._flight_client_options['generic_options'])
+        if disabled:
+            self.assertEqual(generic_options.get('grpc.compression_enabled_algorithms_bitset'), 1)
+        else:
+            self.assertIsNone(generic_options.get('grpc.compression_enabled_algorithms_bitset'))
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_DISABLE_GRPC_COMPRESSION': 'true'})
+    def test_from_env_disable_grpc_compression_true(self):
+        client = InfluxDBClient3.from_env()
+        self.assertGrpcCompressionDisabled(client, True)
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_DISABLE_GRPC_COMPRESSION': 'TrUe'})
+    def test_from_env_disable_grpc_compression_true_mixed_case(self):
+        client = InfluxDBClient3.from_env()
+        self.assertGrpcCompressionDisabled(client, True)
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_DISABLE_GRPC_COMPRESSION': '1'})
+    def test_from_env_disable_grpc_compression_one(self):
+        client = InfluxDBClient3.from_env()
+        self.assertGrpcCompressionDisabled(client, True)
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_DISABLE_GRPC_COMPRESSION': 'false'})
+    def test_from_env_disable_grpc_compression_false(self):
+        client = InfluxDBClient3.from_env()
+        self.assertGrpcCompressionDisabled(client, False)
+
+    @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_DISABLE_GRPC_COMPRESSION': 'anything-else'})
+    def test_from_env_disable_grpc_compression_anything_else_is_false(self):
+        client = InfluxDBClient3.from_env()
+        self.assertGrpcCompressionDisabled(client, False)
+
+    def test_disable_grpc_compression_parameter_true(self):
+        client = InfluxDBClient3(
+            host="localhost",
+            org="my_org",
+            database="my_db",
+            token="my_token",
+            disable_grpc_compression=True
+        )
+        self.assertGrpcCompressionDisabled(client, True)
+
+    def test_disable_grpc_compression_parameter_false(self):
+        client = InfluxDBClient3(
+            host="localhost",
+            org="my_org",
+            database="my_db",
+            token="my_token",
+            disable_grpc_compression=False
+        )
+        self.assertGrpcCompressionDisabled(client, False)
+
+    def test_disable_grpc_compression_default_is_false(self):
+        client = InfluxDBClient3(
+            host="localhost",
+            org="my_org",
+            database="my_db",
+            token="my_token",
+        )
+        self.assertGrpcCompressionDisabled(client, False)
+
     def test_query_with_arrow_error(self):
         f = ErrorFlightServer()
         with InfluxDBClient3(f"http://localhost:{f.port}", "my_org", "my_db", "my_token") as c:

tests/test_influxdb_client_3_integration.py

@@ -4,6 +4,7 @@
 import pytest
 import random
 import string
+import sys
 import time
 import unittest
 
@@ -455,3 +456,115 @@ def retry_cb(args, data, excp):
         self.assertEqual(lp, ErrorResult["rd"].decode('utf-8'))
         self.assertIsNotNone(ErrorResult["rx"])
         self.assertIsInstance(ErrorResult["rx"], Url3TimeoutError)
+
+    def test_disable_grpc_compression(self):
+        """
+        Test that disable_grpc_compression parameter controls query response compression.
+
+        On Python 3.10+: Uses mitmproxy to intercept and verify request+response headers.
+        On Python <3.10: Only verifies data integrity (mitmproxy not available).
+        """
+        # Test cases
+        test_cases = [
+            {
+                'name': 'default',
+                'disable_grpc_compression': None,
+                'expected_req_encoding': 'identity, deflate, gzip',
+                'expected_resp_encoding': 'gzip',
+            },
+            {
+                'name': 'disabled=False',
+                'disable_grpc_compression': False,
+                'expected_req_encoding': 'identity, deflate, gzip',
+                'expected_resp_encoding': 'gzip',
+            },
+            {
+                'name': 'disabled=True',
+                'disable_grpc_compression': True,
+                'expected_req_encoding': 'identity',
+                'expected_resp_encoding': None,
+            },
+        ]
+
+        # Check if mitmproxy is available (Python 3.10+ only)
+        mitmproxy_available = sys.version_info >= (3, 10)
+        proxy = None
+        proxy_url = None
+
+        if mitmproxy_available:
+            try:
+                from tests.util.mitmproxy import MitmproxyServer
+                proxy = MitmproxyServer()
+                proxy.__enter__()
+                proxy_url = proxy.url
+            except ImportError as e:
+                logging.error(f"mitmproxy not available: {e}")
+                mitmproxy_available = False
+
+        try:
+            test_id = time.time_ns()
+            measurement = f'grpc_compression_test_{random_hex(6)}'
+
+            # Write test data points
+            num_points = 10
+            lines = [
+                f'{measurement},type=test value={i}.0,counter={i}i,test_id={test_id}i {test_id + i * 1000000}'
+                for i in range(num_points)
+            ]
+            self.client.write('\n'.join(lines))
+
+            test_query = f"SELECT * FROM \"{measurement}\" WHERE test_id = {test_id} ORDER BY counter"
+
+            # Wait for data to be available
+            result = None
+            start = time.time()
+            while time.time() - start < 10:
+                result = self.client.query(test_query, mode="all")
+                if len(result) >= num_points:
+                    break
+                time.sleep(0.5)
+            self.assertEqual(len(result), num_points, "Data not available after write")
+
+            for tc in test_cases:
+                name = tc['name']
+
+                # Build client kwargs
+                client_kwargs = {
+                    'host': self.host,
+                    'database': self.database,
+                    'token': self.token,
+                    'proxy': proxy_url,
+                    'verify_ssl': False if proxy_url else True,
+                }
+                if tc['disable_grpc_compression'] is not None:
+                    client_kwargs['disable_grpc_compression'] = tc['disable_grpc_compression']
+
+                client = InfluxDBClient3(**client_kwargs)
+                try:
+                    result = client.query(test_query, mode="all")
+                    self.assertEqual(len(result), num_points, f"[{name}] Should return {num_points} rows")
+                finally:
+                    client.close()
+
+                # Verify headers (Python 3.10+ only)
+                if mitmproxy_available and proxy:
+                    req_encoding = proxy.capture.get_last_request_header('grpc-accept-encoding')
+                    resp_encoding = proxy.capture.get_last_response_header('grpc-encoding')
+
+                    print(f"\n[{name}] Request grpc-accept-encoding: {req_encoding}")
+                    print(f"[{name}] Response grpc-encoding: {resp_encoding}")
+
+                    self.assertEqual(req_encoding, tc['expected_req_encoding'],
+                                     f"[{name}] Unexpected request encoding")
+
+                    if tc['expected_resp_encoding']:
+                        self.assertEqual(resp_encoding, tc['expected_resp_encoding'],
+                                         f"[{name}] Unexpected response encoding")
+                    else:
+                        self.assertTrue(resp_encoding is None or resp_encoding == 'identity',
+                                        f"[{name}] Expected no compression, got: {resp_encoding}")
+
+                    proxy.capture.clear()
+        finally:
+            if proxy:
+                proxy.__exit__(None, None, None)