feat: support env

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -10,7 +10,7 @@
 from influxdb_client_3.write_client import InfluxDBClient as _InfluxDBClient, WriteOptions, Point
 from influxdb_client_3.write_client.client.exceptions import InfluxDBError
 from influxdb_client_3.write_client.client.write_api import WriteApi as _WriteApi, SYNCHRONOUS, ASYNCHRONOUS, \
-    PointSettings, WriteType
+    PointSettings, WriteType, DefaultWriteOptions
 from influxdb_client_3.write_client.domain.write_precision import WritePrecision
 
 polars = importlib.util.find_spec("polars") is not None
@@ -57,6 +57,7 @@ def file_parser_options(**kwargs):
 INFLUX_ORG = "INFLUX_ORG"
 INFLUX_PRECISION = "INFLUX_PRECISION"
 INFLUX_AUTH_SCHEME = "INFLUX_AUTH_SCHEME"
+INFLUX_GZIP_THRESHOLD = "INFLUX_GZIP_THRESHOLD"
 
 
 def from_env(**kwargs: Any) -> 'InfluxDBClient3':
@@ -92,16 +93,20 @@ def from_env(**kwargs: Any) -> 'InfluxDBClient3':
     if missing_vars:
         raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}")
 
-    org = os.getenv(INFLUX_ORG, "default")
+    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))
+
+    if os.getenv(INFLUX_PRECISION) is not None:
+        write_options.write_precision = os.getenv(INFLUX_PRECISION)
+
+    write_client_option = {'write_options': write_options}
 
     if os.getenv(INFLUX_AUTH_SCHEME) is not None:
         kwargs['auth_scheme'] = os.getenv(INFLUX_AUTH_SCHEME)
-    
-    write_client_option = None
-    if os.getenv(INFLUX_PRECISION) is not None:
-        write_client_option = default_client_options(
-            write_options=WriteOptions(write_type=WriteType.synchronous, write_precision=os.getenv(INFLUX_PRECISION))
-        )
+
+    org = os.getenv(INFLUX_ORG, "default")
 
     return InfluxDBClient3(
         host=required_vars[INFLUX_HOST],
@@ -202,8 +207,26 @@ def __init__(
         self._org = org if org is not None else "default"
         self._database = database
         self._token = token
-        self._write_client_options = write_client_options if write_client_options is not None \
-            else default_client_options(write_options=WriteOptions(write_type=WriteType.synchronous, write_precision=WritePrecision.NS))
+
+        write_type = DefaultWriteOptions['write_type']
+        write_precision = DefaultWriteOptions['write_precision']
+        gzip_threshold = DefaultWriteOptions['gzip_threshold']
+        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')
+            gzip_threshold = getattr(write_opts, 'gzip_threshold')
+
+        write_options = WriteOptions(
+            write_type=write_type,
+            write_precision=write_precision,
+            gzip_threshold=gzip_threshold,
+        )
+
+        self._write_client_options = {
+            "write_options": write_options,
+            **(write_client_options or {})
+        }
 
         # Parse the host input
         parsed_url = urllib.parse.urlparse(host)

influxdb_client_3/write_client/client/write_api.py

@@ -39,6 +39,13 @@ class WriteType(Enum):
     synchronous = 3
 
 
+DefaultWriteOptions = {
+    'write_type': WriteType.synchronous,
+    'write_precision': WritePrecision.NS,
+    'gzip_threshold': 1000
+}
+
+
 class WriteOptions(object):
     """Write configuration."""
 
@@ -52,6 +59,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
                  exponential_base=2,
                  max_close_wait=300_000,
                  write_precision=DEFAULT_WRITE_PRECISION,
+                 gzip_threshold=1000,
                  write_scheduler=ThreadPoolScheduler(max_workers=1)) -> None:
         """
         Create write api configuration.
@@ -83,6 +91,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
         self.write_scheduler = write_scheduler
         self.max_close_wait = max_close_wait
         self.write_precision = write_precision
+        self.gzip_threshold = gzip_threshold
 
     def to_retry_strategy(self, **kwargs):
         """

tests/test_influxdb_client_3.py

@@ -78,13 +78,13 @@ def test_default_client(self):
         expected_precision = WritePrecision.NS
         expected_write_type = WriteType.synchronous
 
-        def verify_client_write_options(client):
-            write_options = client._write_client_options.get('write_options')
+        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(client._write_api._write_options.write_precision, expected_precision)
-            self.assertEqual(client._write_api._write_options.write_type, expected_write_type)
+            self.assertEqual(c._write_api._write_options.write_precision, expected_precision)
+            self.assertEqual(c._write_api._write_options.write_type, expected_write_type)
 
         env_client = from_env()
         verify_client_write_options(env_client)
@@ -94,15 +94,20 @@ def verify_client_write_options(client):
 
 
     @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
-                               'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org', 'INFLUX_PRECISION': WritePrecision.MS})
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org', 'INFLUX_PRECISION': WritePrecision.MS,
+                               'INFLUX_GZIP_THRESHOLD': '2000'})
     def test_from_env_all_env_vars_set(self):
         client = from_env()
         self.assertIsInstance(client, InfluxDBClient3)
         self.assertEqual(client._client.url, "https://localhost:443")
         self.assertEqual(client._database, "test_db")
         self.assertEqual(client._org, "test_org")
         self.assertEqual(client._token, "test_token")
-        self.assertEqual(client._write_client_options.get("write_options").write_precision, WritePrecision.MS)
+
+        write_options = client._write_client_options.get("write_options")
+        self.assertEqual(write_options.write_precision, WritePrecision.MS)
+        self.assertEqual(write_options.gzip_threshold, 2000)
+
 
     @patch.dict('os.environ', {'INFLUX_HOST': "", 'INFLUX_TOKEN': "",
                                'INFLUX_DATABASE': "", 'INFLUX_ORG': ""})