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
+    PointSettings, WriteType
 from influxdb_client_3.write_client.domain.write_precision import WritePrecision
 
 polars = importlib.util.find_spec("polars") is not None
@@ -55,9 +55,11 @@ def file_parser_options(**kwargs):
 INFLUX_TOKEN = "INFLUX_TOKEN"
 INFLUX_DATABASE = "INFLUX_DATABASE"
 INFLUX_ORG = "INFLUX_ORG"
+INFLUX_PRECISION = "INFLUX_PRECISION"
 
 
 def from_env(**kwargs: Any) -> 'InfluxDBClient3':
+    
     """
     Create an instance of `InfluxDBClient3` using environment variables for configuration.
 
@@ -78,7 +80,7 @@ def from_env(**kwargs: Any) -> 'InfluxDBClient3':
                    flight_client_options, SSL settings, etc.
     :return: An initialized `InfluxDBClient3` instance.
     :raises ValueError: If any required environment variables are not set.
-    """
+    """ 
     required_vars = {
         INFLUX_HOST: os.getenv(INFLUX_HOST),
         INFLUX_TOKEN: os.getenv(INFLUX_TOKEN),
@@ -91,10 +93,17 @@ def from_env(**kwargs: Any) -> 'InfluxDBClient3':
 
     org = os.getenv(INFLUX_ORG, "default")
 
+    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))
+        )
+
     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,
         **kwargs
     )
@@ -190,7 +199,7 @@ def __init__(
         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=SYNCHRONOUS)
+            else default_client_options(write_options=WriteOptions(write_type=WriteType.synchronous, write_precision=WritePrecision.NS))
 
         # Parse the host input
         parsed_url = urllib.parse.urlparse(host)

influxdb_client_3/write_client/client/write_api.py

@@ -51,6 +51,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
                  max_retry_time=180_000,
                  exponential_base=2,
                  max_close_wait=300_000,
+                 write_precision=DEFAULT_WRITE_PRECISION,
                  write_scheduler=ThreadPoolScheduler(max_workers=1)) -> None:
         """
         Create write api configuration.
@@ -66,7 +67,8 @@ def __init__(self, write_type: WriteType = WriteType.batching,
         :param max_retry_delay: the maximum delay between each retry attempt in milliseconds
         :param max_retry_time: total timeout for all retry attempts in milliseconds, if 0 retry is disabled
         :param exponential_base: base for the exponential retry delay
-        :parama max_close_wait: the maximum time to wait for writes to be flushed if close() is called
+        :param max_close_wait: the maximum time to wait for writes to be flushed if close() is called
+        :param write_precision: the time precision for the data written to InfluxDB.
         :param write_scheduler:
         """
         self.write_type = write_type
@@ -80,6 +82,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
         self.exponential_base = exponential_base
         self.write_scheduler = write_scheduler
         self.max_close_wait = max_close_wait
+        self.write_precision = write_precision
 
     def to_retry_strategy(self, **kwargs):
         """
@@ -290,7 +293,7 @@ def write(self, bucket: str, org: str = None,
                   str, Iterable['str'], Point, Iterable['Point'], dict, Iterable['dict'], bytes, Iterable['bytes'],
                   Observable, NamedTuple, Iterable['NamedTuple'], 'dataclass', Iterable['dataclass']
               ] = None,
-              write_precision: WritePrecision = DEFAULT_WRITE_PRECISION, **kwargs) -> Any:
+              write_precision: WritePrecision = None, **kwargs) -> Any:
         """
         Write time-series data into InfluxDB.
 
@@ -360,7 +363,10 @@ def write(self, bucket: str, org: str = None,
         org = get_org_query_param(org=org, client=self._influxdb_client)
 
         self._append_default_tags(record)
-
+        
+        if write_precision is None:
+            write_precision = self._write_options.write_precision
+        
         if self._write_options.write_type is WriteType.batching:
             return self._write_batching(bucket, org, record,
                                         write_precision, **kwargs)
@@ -443,8 +449,11 @@ def __del__(self):
         pass
 
     def _write_batching(self, bucket, org, data,
-                        precision=DEFAULT_WRITE_PRECISION,
+                        precision=None,
                         **kwargs):
+        if precision is None:
+            precision = self._write_options.write_precision
+            
         if isinstance(data, bytes):
             _key = _BatchItemKey(bucket, org, precision)
             self._subject.on_next(_BatchItem(key=_key, data=data))
@@ -454,7 +463,8 @@ def _write_batching(self, bucket, org, data,
                                  precision, **kwargs)
 
         elif isinstance(data, Point):
-            self._write_batching(bucket, org, data.to_line_protocol(), data.write_precision, **kwargs)
+            write_precision = data.write_precision if data.write_precision is not None else precision
+            self._write_batching(bucket, org, data.to_line_protocol(), write_precision, **kwargs)
 
         elif isinstance(data, dict):
             self._write_batching(bucket, org, Point.from_dict(data, write_precision=precision, **kwargs),

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, write_client_options
+from influxdb_client_3 import InfluxDBClient3, from_env, WritePrecision, WriteType
 from tests.util import asyncio_run
 from tests.util.mocks import ConstantFlightServer, ConstantData
 
@@ -74,22 +74,35 @@ async def test_query_async(self):
             assert {'data': 'sql_query', 'reference': query, 'value': -1.0} in result_list
             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
+
+        def verify_client_write_options(client):
+            write_options = client._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)
+
+        env_client = from_env()
+        verify_client_write_options(env_client)
+
+        default_client = InfluxDBClient3()
+        verify_client_write_options(default_client)
+
+
     @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
-                               'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org'})
+                               'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org', 'INFLUX_PRECISION': WritePrecision.MS})
     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")
-
-    def test_from_env_with_kargs(self):
-        client = from_env(
-            write_client_options=write_client_options(batch_size=10000),
-        )
-        self.assertIsInstance(client, InfluxDBClient3)
-        self.assertEqual(client._write_client_options['batch_size'], 10000)
+        self.assertEqual(client._write_client_options.get("write_options").write_precision, WritePrecision.MS)
 
     @patch.dict('os.environ', {'INFLUX_HOST': "", 'INFLUX_TOKEN': "",
                                'INFLUX_DATABASE': "", 'INFLUX_ORG': ""})