feat: improve error handling

Author: NguyenHoangSon96

influxdb_client_3/__init__.py

@@ -4,7 +4,9 @@
 from typing import Any
 
 import pyarrow as pa
+from pyarrow import ArrowException
 
+from influxdb_client_3.influxdb_client_error import InfluxdbClientQueryError
 from influxdb_client_3.query.query_api import QueryApi as _QueryApi, QueryApiOptionsBuilder
 from influxdb_client_3.read_file import UploadFile
 from influxdb_client_3.write_client import InfluxDBClient as _InfluxDBClient, WriteOptions, Point
@@ -400,8 +402,8 @@ def query(self, query: str, language: str = "sql", mode: str = "all", database:
 
         try:
             return self._query_api.query(query=query, language=language, mode=mode, database=database, **kwargs)
-        except InfluxDBError as e:
-            raise e
+        except ArrowException as e:
+            raise InfluxdbClientQueryError(f"Error while executing query: {e}")
 
     async def query_async(self, query: str, language: str = "sql", mode: str = "all", database: str = None, **kwargs):
         """Query data from InfluxDB asynchronously.
@@ -428,13 +430,13 @@ async def query_async(self, query: str, language: str = "sql", mode: str = "all"
             database = self._database
 
         try:
-            return await self._query_api.query_async(query=query,
+             return await self._query_api.query_async(query=query,
                                                      language=language,
                                                      mode=mode,
                                                      database=database,
                                                      **kwargs)
-        except InfluxDBError as e:
-            raise e
+        except ArrowException as e:
+            raise InfluxdbClientQueryError(f"Error while executing query: {e}")
 
     def close(self):
         """Close the client and clean up resources."""

influxdb_client_3/influxdb_client_error.py

@@ -0,0 +1,12 @@
+class InfluxDBClientError(Exception):
+    pass
+
+
+class InfluxdbClientQueryError(InfluxDBClientError):
+    def __init__(self, error_message, *args, **kwargs):
+        super().__init__(error_message, *args, **kwargs)
+        self.message = error_message
+
+
+class InfluxdbClientWriteError(InfluxDBClientError):
+    pass

influxdb_client_3/query/query_api.py

@@ -162,14 +162,12 @@ def query(self, query: str, language: str, mode: str, database: str, **kwargs):
                                    It should be a ``dictionary`` of key-value pairs.
         :return: The query result in the specified mode.
         """
-        try:
-            ticket, _options = self._prepare_query(query, language, database, **kwargs)
+        ticket, _options = self._prepare_query(query, language, database, **kwargs)
 
-            flight_reader = self._do_get(ticket, _options)
+        flight_reader = self._do_get(ticket, _options)
+
+        return self._translate_stream_reader(flight_reader, mode)
 
-            return self._translate_stream_reader(flight_reader, mode)
-        except Exception as e:
-            raise e
 
     async def query_async(self, query: str, language: str, mode: str, database: str, **kwargs):
         """Query data from InfluxDB asynchronously.
@@ -187,16 +185,13 @@ async def query_async(self, query: str, language: str, mode: str, database: str,
                            It should be a ``dictionary`` of key-value pairs.
         :return: The query result in the specified mode.
         """
-        try:
-            ticket, options = self._prepare_query(query, language, database, **kwargs)
-            loop = asyncio.get_running_loop()
-            _flight_reader = await loop.run_in_executor(None,
-                                                        self._flight_client.do_get, ticket, options)
-            return await loop.run_in_executor(None, self._translate_stream_reader,
+        ticket, options = self._prepare_query(query, language, database, **kwargs)
+        loop = asyncio.get_running_loop()
+        _flight_reader = await loop.run_in_executor(None,
+                                                    self._flight_client.do_get, ticket, options)
+        return await loop.run_in_executor(None, self._translate_stream_reader,
                                               _flight_reader,
                                               mode)
-        except Exception as e:
-            raise e
 
     def _translate_stream_reader(self, reader: FlightStreamReader, mode: str):
         from influxdb_client_3 import polars as has_polars

tests/test_influxdb_client_3.py

@@ -2,8 +2,9 @@
 from unittest.mock import patch
 
 from influxdb_client_3 import InfluxDBClient3, WritePrecision, DefaultWriteOptions, Point, WriteOptions, WriteType
+from influxdb_client_3.influxdb_client_error import InfluxdbClientQueryError
 from tests.util import asyncio_run
-from tests.util.mocks import ConstantFlightServer, ConstantData
+from tests.util.mocks import ConstantFlightServer, ConstantData, ErrorFlightServer
 
 
 class TestInfluxDBClient3(unittest.TestCase):
@@ -142,6 +143,20 @@ def test_parse_invalid_write_precision(self):
             InfluxDBClient3.from_env()
         self.assertIn("Invalid precision value: invalid_value", str(context.exception))
 
+    def test_query_with_arrow_error(self):
+        f = ErrorFlightServer()
+        with InfluxDBClient3(f"http://localhost:{f.port}", "my_org", "my_db", "my_token") as c:
+            with self.assertRaises(InfluxdbClientQueryError) as err:
+                c.query("SELECT * FROM my_data")
+            self.assertIn("Error while executing query", str(err.exception))
+
+    @asyncio_run
+    async def test_async_query_with_arrow_error(self):
+        f = ErrorFlightServer()
+        with InfluxDBClient3(f"http://localhost:{f.port}", "my_org", "my_db", "my_token") as c:
+            with self.assertRaises(InfluxdbClientQueryError) as err:
+                await c.query_async("SELECT * FROM my_data")
+            self.assertIn("Error while executing query", str(err.exception))
 
 if __name__ == '__main__':
     unittest.main()

tests/util/mocks.py

@@ -5,9 +5,8 @@
 from pyarrow import (
     array,
     Table,
-    concat_tables
+    concat_tables, ArrowException
 )
-
 from pyarrow.flight import (
     FlightServerBase,
     RecordBatchStream,
@@ -158,3 +157,7 @@ def number_batches(table):
         for idx, batch in enumerate(table.to_batches()):
             buf = struct.pack('<i', idx)
             yield batch, buf
+
+class ErrorFlightServer(FlightServerBase):
+    def do_get(self, context, ticket):
+        raise ArrowException