chore: minor refactoring. Sync timeout types in __init__. Extra tests.

Author: karel-rehor

influxdb_client_3/__init__.py

@@ -25,6 +25,8 @@
 INFLUX_AUTH_SCHEME = "INFLUX_AUTH_SCHEME"
 INFLUX_GZIP_THRESHOLD = "INFLUX_GZIP_THRESHOLD"
 INFLUX_WRITE_NO_SYNC = "INFLUX_WRITE_NO_SYNC"
+INFLUX_WRITE_TIMEOUT = "INFLUX_WRITE_TIMEOUT"
+INFLUX_QUERY_TIMEOUT = "INFLUX_QUERY_TIMEOUT"
 
 
 def write_client_options(**kwargs):
@@ -97,7 +99,7 @@ def _merge_options(defaults, exclude_keys=None, custom=None):
     return _deep_merge(defaults, {key: value for key, value in custom.items() if key not in exclude_keys})
 
 
-def _parse_precision(precision):
+def _parse_precision(precision: str) -> WritePrecision:
     """
     Parses the precision value and ensures it is valid.
 
@@ -124,7 +126,7 @@ def _parse_precision(precision):
     raise ValueError(f"Invalid precision value: {precision}")
 
 
-def _parse_gzip_threshold(threshold):
+def _parse_gzip_threshold(threshold: str) -> int:
     """
     Parses and validates the provided threshold value.
 
@@ -148,7 +150,7 @@ def _parse_gzip_threshold(threshold):
     return threshold
 
 
-def _parse_write_no_sync(write_no_sync):
+def _parse_write_no_sync(write_no_sync: str):
     """
     Parses and validates the provided write no sync value.
 
@@ -163,6 +165,16 @@ def _parse_write_no_sync(write_no_sync):
     return write_no_sync.strip().lower() in ['true', '1', 't', 'y', 'yes']
 
 
+def _parse_timeout(to: str) -> int:
+    try:
+        timeout = int(to)
+    except (TypeError, ValueError):
+        raise ValueError(f"Invalid timeout value: {to}. Must be a number.")
+    if timeout < 0:
+        raise ValueError(f"Invalid timeout value: {to}. Must be non-negative.")
+    return timeout
+
+
 class InfluxDBClient3:
     def __init__(
             self,
@@ -211,24 +223,30 @@ def __init__(
                               (defaults to false, don't set to true when talking to InfluxDB 2)
         :key str username: ``username`` to authenticate via username and password credentials to the InfluxDB 2.x
         :key str password: ``password`` to authenticate via username and password credentials to the InfluxDB 2.x
-        :key str query_timeout: float value used to set the client query API timeout in seconds.
+        :key str query_timeout: int value used to set the client query API timeout in milliseconds.
+        :key str write_timeout: int value used to set the client write API timeout in milliseconds.
         :key list[str] profilers: list of enabled Flux profilers
         """
         self._org = org if org is not None else "default"
         self._database = database
         self._token = token
+        kw_keys = kwargs.keys()
 
         write_type = DefaultWriteOptions.write_type.value
         write_precision = DefaultWriteOptions.write_precision.value
         write_no_sync = DefaultWriteOptions.no_sync.value
         write_timeout = DefaultWriteOptions.timeout.value
+
         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_type)
             write_precision = getattr(write_opts, 'write_precision', write_precision)
             write_no_sync = getattr(write_opts, 'no_sync', write_no_sync)
             write_timeout = getattr(write_opts, 'timeout', write_timeout)
 
+        if kw_keys.__contains__('write_timeout'):
+            write_timeout = kwargs.get('write_timeout')
+
         write_options = WriteOptions(
             write_type=write_type,
             write_precision=write_precision,
@@ -269,15 +287,15 @@ def __init__(
             connection_string = f"grpc+tcp://{hostname}:{port}"
 
         q_opts_builder = QueryApiOptionsBuilder()
-        kw_keys = kwargs.keys()
         if kw_keys.__contains__('ssl_ca_cert'):
             q_opts_builder.root_certs(kwargs.get('ssl_ca_cert', None))
         if kw_keys.__contains__('verify_ssl'):
             q_opts_builder.tls_verify(kwargs.get('verify_ssl', True))
         if kw_keys.__contains__('proxy'):
             q_opts_builder.proxy(kwargs.get('proxy', None))
         if kw_keys.__contains__('query_timeout'):
-            q_opts_builder.timeout(kwargs.get('query_timeout', None))
+            query_timeout_float = float(kwargs.get('query_timeout'))
+            q_opts_builder.timeout(query_timeout_float / 1000.0)
         self._query_api = _QueryApi(connection_string=connection_string, token=token,
                                     flight_client_options=flight_client_options,
                                     proxy=kwargs.get("proxy", None), options=q_opts_builder.build())
@@ -325,6 +343,15 @@ def from_env(cls, **kwargs: Any) -> 'InfluxDBClient3':
         if precision is not None:
             write_options.write_precision = _parse_precision(precision)
 
+        write_timeout = os.getenv(INFLUX_WRITE_TIMEOUT)
+        if write_timeout is not None:
+            # N.B. write_options value has precedent over kwargs['write_timeout'] above
+            write_options.timeout = _parse_timeout(write_timeout)
+
+        query_timeout = os.getenv(INFLUX_QUERY_TIMEOUT)
+        if query_timeout is not None:
+            kwargs['query_timeout'] = _parse_timeout(query_timeout)
+
         write_client_option = {'write_options': write_options}
 
         if os.getenv(INFLUX_AUTH_SCHEME) is not None:

influxdb_client_3/query/query_api.py

@@ -18,6 +18,7 @@ class QueryApiOptions(object):
     tls_verify (bool): whether to verify SSL certificates or not
     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
     """
     _DEFAULT_TIMEOUT = 300.0
     tls_root_certs: bytes = None
@@ -26,7 +27,11 @@ class QueryApiOptions(object):
     flight_client_options: dict = None
     timeout: float = None
 
-    def __init__(self, root_certs_path, verify, proxy, flight_client_options, timeout):
+    def __init__(self, root_certs_path: str,
+                 verify: bool,
+                 proxy: str,
+                 flight_client_options: dict,
+                 timeout: float = _DEFAULT_TIMEOUT):
         """
         Initialize a set of QueryApiOptions
 
@@ -35,6 +40,7 @@ def __init__(self, root_certs_path, verify, proxy, flight_client_options, timeou
         :param proxy: URL of a proxy server, if required.
         :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.
         """
         if root_certs_path:
             self.tls_root_certs = self._read_certs(root_certs_path)
@@ -43,7 +49,7 @@ def __init__(self, root_certs_path, verify, proxy, flight_client_options, timeou
         self.flight_client_options = flight_client_options
         self.timeout = timeout
 
-    def _read_certs(self, path):
+    def _read_certs(self, path: str) -> bytes:
         with open(path, "rb") as certs_file:
             return certs_file.read()
 
@@ -64,33 +70,33 @@ class QueryApiOptionsBuilder(object):
 
         client = QueryApi(connection, token, None, None, options)
     """
-    _root_certs_path = None
-    _tls_verify = True
-    _proxy = None
-    _flight_client_options = None
-    _timeout = None
+    _root_certs_path: str = None
+    _tls_verify: bool = True
+    _proxy: str = None
+    _flight_client_options: dict = None
+    _timeout: float = None
 
-    def root_certs(self, path):
+    def root_certs(self, path: str):
         self._root_certs_path = path
         return self
 
-    def tls_verify(self, verify):
+    def tls_verify(self, verify: bool):
         self._tls_verify = verify
         return self
 
-    def proxy(self, proxy):
+    def proxy(self, proxy: str):
         self._proxy = proxy
         return self
 
-    def flight_client_options(self, flight_client_options):
+    def flight_client_options(self, flight_client_options: dict):
         self._flight_client_options = flight_client_options
         return self
 
-    def timeout(self, timeout):
+    def timeout(self, timeout: float):
         self._timeout = timeout
         return self
 
-    def build(self):
+    def build(self) -> QueryApiOptions:
         """Build a QueryApiOptions object with previously set values"""
         return QueryApiOptions(
             root_certs_path=self._root_certs_path,

influxdb_client_3/write_client/client/write_api.py

@@ -82,6 +82,7 @@ def __init__(self, write_type: WriteType = WriteType.batching,
         :param max_close_wait: the maximum time to wait for writes to be flushed if close() is called
         :param write_precision: precision to use when writing points to InfluxDB
         :param no_sync: skip waiting for WAL persistence on write
+        :param timeout: timeout to use when writing to the database in milliseconds. Default is 10_000
         :param write_scheduler:
         """
         self.write_type = write_type
@@ -397,9 +398,7 @@ def write(self, bucket: str, org: str = None,
 
         def write_payload(payload):
             final_string = b'\n'.join(payload[1])
-            result = self._post_write(_async_req, bucket, org, final_string, payload[0], no_sync)
-            # return self._post_write(_async_req, bucket, org, final_string, payload[0], no_sync)
-            return result
+            return self._post_write(_async_req, bucket, org, final_string, payload[0], no_sync)
 
         results = list(map(write_payload, payloads.items()))
         if not _async_req:

tests/test_influxdb_client_3.py

@@ -187,7 +187,8 @@ def verify_client_write_options(c):
     @patch.dict('os.environ', {'INFLUX_HOST': 'localhost', 'INFLUX_TOKEN': 'test_token',
                                'INFLUX_DATABASE': 'test_db', 'INFLUX_ORG': 'test_org',
                                'INFLUX_PRECISION': WritePrecision.MS, 'INFLUX_AUTH_SCHEME': 'custom_scheme',
-                               'INFLUX_GZIP_THRESHOLD': '2000', 'INFLUX_WRITE_NO_SYNC': 'true'})
+                               'INFLUX_GZIP_THRESHOLD': '2000', 'INFLUX_WRITE_NO_SYNC': 'true',
+                               'INFLUX_WRITE_TIMEOUT': '1234', 'INFLUX_QUERY_TIMEOUT': '5678'})
     def test_from_env_all_env_vars_set(self):
         client = InfluxDBClient3.from_env()
         self.assertIsInstance(client, InfluxDBClient3)
@@ -201,6 +202,8 @@ def test_from_env_all_env_vars_set(self):
         write_options = client._write_client_options.get("write_options")
         self.assertEqual(write_options.write_precision, WritePrecision.MS)
         self.assertEqual(write_options.no_sync, True)
+        self.assertEqual(1234, write_options.timeout)
+        self.assertEqual(5.678, client._query_api._default_timeout)
 
         client._write_api._point_settings = {}
 

tests/test_influxdb_client_3_integration.py

@@ -251,6 +251,22 @@ def test_get_server_version(self):
         version = self.client.get_server_version()
         assert version is not None
 
+    def test_write_timeout(self):
+        with pytest.raises(Url3TimeoutError):
+            InfluxDBClient3(
+                host=self.host,
+                database=self.database,
+                token=self.token,
+                write_timeout=30,
+                write_client_options=write_client_options(
+                    write_options=WriteOptions(
+                        max_retry_time=0,
+                        timeout=20,
+                        write_type=WriteType.synchronous
+                    )
+                )
+            ).write("test_write_timeout,location=harfa fVal=3.14,iVal=42i")
+
     def test_write_timeout_sync(self):
 
         with pytest.raises(Url3TimeoutError):
@@ -376,7 +392,7 @@ def test_query_timeout(self):
             host=self.host,
             token=self.token,
             database=self.database,
-            query_timeout=0.001,
+            query_timeout=1,
         )
 
         with self.assertRaisesRegex(InfluxDB3ClientQueryError, ".*Deadline Exceeded.*"):
@@ -387,7 +403,7 @@ def test_query_timeout_per_call_override(self):
             host=self.host,
             token=self.token,
             database=self.database,
-            query_timeout=3.0,
+            query_timeout=3,
         )
 
         with self.assertRaisesRegex(InfluxDB3ClientQueryError, ".*Deadline Exceeded.*"):

tests/test_write_local_server.py

@@ -134,7 +134,7 @@ def test_write_with_no_sync_true_and_gzip(self, httpserver: HTTPServer):
             query_string={"org": "ORG", "db": "DB", "precision": "microsecond", "no_sync": "true"},
             headers={"Content-Encoding": "gzip"}, ))
 
-    def test_write_with_timeout(self, httpserver: HTTPServer):
+    def test_write_with_timeout_in_write_options(self, httpserver: HTTPServer):
         self.delay_response(httpserver, 0.5)
 
         with pytest.raises(urllib3_TimeoutError):
@@ -150,3 +150,20 @@ def test_write_with_timeout(self, httpserver: HTTPServer):
                 ),
                 enable_gzip=True
             ).write(self.SAMPLE_RECORD)
+
+    def test_write_with_write_timeout(self, httpserver: HTTPServer):
+        self.delay_response(httpserver, 0.5)
+
+        with pytest.raises(urllib3_TimeoutError):
+            InfluxDBClient3(
+                host=(httpserver.url_for("/")), org="ORG", database="DB", token="TOKEN",
+                write_timeout=30,
+                write_client_options=write_client_options(
+                    write_options=WriteOptions(
+                        write_type=WriteType.synchronous,
+                        write_precision=WritePrecision.US,
+                        no_sync=True,
+                    )
+                ),
+                enable_gzip=True
+            ).write(self.SAMPLE_RECORD)