feat: improve error handling (#140)

NguyenHoangSon96 · web-flow · commit 39bbea99205c · 2025-05-27T21:00:19.000+07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,7 +7,8 @@
 ### Features
 
 1. [#130](https://github.com/InfluxCommunity/influxdb3-python/pull/130): Remove org parameters from the example code because It is not mandatory in Influxdb3
-2. [#139](https://github.com/InfluxCommunity/influxdb3-python/pull/139): Supports environment variables with the same name like other clients   
+2. [#139](https://github.com/InfluxCommunity/influxdb3-python/pull/139): Supports environment variables with the same name as other clients
+3. [#140](https://github.com/InfluxCommunity/influxdb3-python/pull/140): Query api will throw `InfluxdbClientQueryError` when receiving `ArrowException` from gRPC servers
 
 ## 0.12.0 [2025-03-26] 
 
diff --git a/Examples/handle_query_error.py b/Examples/handle_query_error.py
@@ -0,0 +1,33 @@
+"""
+Demonstrates handling error when querying InfluxDB.
+"""
+import logging
+from config import Config
+from influxdb_client_3.exceptions import InfluxDB3ClientQueryError
+
+import influxdb_client_3 as InfluxDBClient3
+
+
+def main() -> None:
+    """
+    Main function
+    :return:
+    """
+    config = Config()
+    logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+
+    client = InfluxDBClient3.InfluxDBClient3(
+        host=config.host,
+        token=config.token,
+        database=config.database
+    )
+
+    try:
+        # Select from a bucket that doesn't exist
+        client.query("Select a from cpu11")
+    except InfluxDB3ClientQueryError as e:
+        logging.log(logging.ERROR, e.message)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/influxdb_client_3/__init__.py b/influxdb_client_3/__init__.py
@@ -4,11 +4,13 @@
 from typing import Any
 
 import pyarrow as pa
+from pyarrow import ArrowException
 
+from influxdb_client_3.exceptions import InfluxDB3ClientQueryError
 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
-from influxdb_client_3.write_client.client.exceptions import InfluxDBError
+from influxdb_client_3.exceptions import InfluxDBError
 from influxdb_client_3.write_client.client.write_api import WriteApi as _WriteApi, SYNCHRONOUS, ASYNCHRONOUS, \
     PointSettings, DefaultWriteOptions, WriteType
 from influxdb_client_3.write_client.domain.write_precision import WritePrecision
@@ -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 InfluxDB3ClientQueryError(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.
@@ -433,8 +435,8 @@ async def query_async(self, query: str, language: str = "sql", mode: str = "all"
                                                      mode=mode,
                                                      database=database,
                                                      **kwargs)
-        except InfluxDBError as e:
-            raise e
+        except ArrowException as e:
+            raise InfluxDB3ClientQueryError(f"Error while executing query: {e}")
 
     def close(self):
         """Close the client and clean up resources."""
diff --git a/influxdb_client_3/exceptions/__init__.py b/influxdb_client_3/exceptions/__init__.py
@@ -0,0 +1,3 @@
+# flake8: noqa
+
+from .exceptions import InfluxDB3ClientQueryError, InfluxDBError, InfluxDB3ClientError
diff --git a/influxdb_client_3/exceptions/exceptions.py b/influxdb_client_3/exceptions/exceptions.py
@@ -4,10 +4,43 @@
 
 from urllib3 import HTTPResponse
 
-logger = logging.getLogger('influxdb_client_3.write_client.client.exceptions')
+logger = logging.getLogger('influxdb_client_3.exceptions')
 
 
-class InfluxDBError(Exception):
+class InfluxDB3ClientError(Exception):
+    """
+    Exception raised for errors in the InfluxDB client operations.
+
+    Represents errors that occur during interactions with the InfluxDB
+    database client. This exception is a general base class for more
+    specific client-related failures and is typically used to signal issues
+    such as invalid queries, connection failures, or API misusage.
+    """
+    pass
+
+
+# This error is for all query operations
+class InfluxDB3ClientQueryError(InfluxDB3ClientError):
+    """
+    Represents an error that occurs when querying an InfluxDB client.
+
+    This class is specifically designed to handle errors originating from
+    client queries to an InfluxDB database. It extends the general
+    `InfluxDBClientError`, allowing more precise identification and
+    handling of query-related issues.
+
+    :ivar message: Contains the specific error message describing the
+        query error.
+    :type message: str
+    """
+
+    def __init__(self, error_message, *args, **kwargs):
+        super().__init__(error_message, *args, **kwargs)
+        self.message = error_message
+
+
+# This error is for all write operations
+class InfluxDBError(InfluxDB3ClientError):
     """Raised when a server error occurs."""
 
     def __init__(self, response: HTTPResponse = None, message: str = None):
diff --git a/influxdb_client_3/query/query_api.py b/influxdb_client_3/query/query_api.py
@@ -162,14 +162,11 @@ 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)
-        except Exception as e:
-            raise e
+        return self._translate_stream_reader(flight_reader, mode)
 
     async def query_async(self, query: str, language: str, mode: str, database: str, **kwargs):
         """Query data from InfluxDB asynchronously.
@@ -187,16 +184,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,
-                                              _flight_reader,
-                                              mode)
-        except Exception as e:
-            raise e
+        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)
 
     def _translate_stream_reader(self, reader: FlightStreamReader, mode: str):
         from influxdb_client_3 import polars as has_polars
diff --git a/influxdb_client_3/write_client/client/util/multiprocessing_helper.py b/influxdb_client_3/write_client/client/util/multiprocessing_helper.py
@@ -8,7 +8,7 @@
 import multiprocessing
 
 from influxdb_client_3.write_client import InfluxDBClient, WriteOptions
-from influxdb_client_3.write_client.client.exceptions import InfluxDBError
+from influxdb_client_3.exceptions import InfluxDBError
 
 logger = logging.getLogger('influxdb_client.client.util.multiprocessing_helper')
 
diff --git a/influxdb_client_3/write_client/client/write/retry.py b/influxdb_client_3/write_client/client/write/retry.py
@@ -9,7 +9,7 @@
 from urllib3 import Retry
 from urllib3.exceptions import MaxRetryError, ResponseError
 
-from influxdb_client_3.write_client.client.exceptions import InfluxDBError
+from influxdb_client_3.exceptions import InfluxDBError
 
 logger = logging.getLogger('influxdb_client.client.write.retry')
 
diff --git a/influxdb_client_3/write_client/rest.py b/influxdb_client_3/write_client/rest.py
@@ -5,7 +5,7 @@
 import logging
 from typing import Dict
 
-from influxdb_client_3.write_client.client.exceptions import InfluxDBError
+from influxdb_client_3.exceptions import InfluxDBError
 from influxdb_client_3.write_client.configuration import Configuration
 
 _UTF_8_encoding = 'utf-8'
diff --git a/tests/test_api_client.py b/tests/test_api_client.py
@@ -6,7 +6,7 @@
 
 from influxdb_client_3.write_client._sync.api_client import ApiClient
 from influxdb_client_3.write_client.configuration import Configuration
-from influxdb_client_3.write_client.client.exceptions import InfluxDBError
+from influxdb_client_3.exceptions import InfluxDBError
 from influxdb_client_3.write_client.service import WriteService
 from influxdb_client_3.version import VERSION
 
diff --git a/tests/test_influxdb_client_3.py b/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.exceptions import InfluxDB3ClientQueryError
 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):
@@ -83,11 +84,29 @@ def test_write_api_custom_options_no_error(self):
             self.assertTrue(True)
         except Exception as e:
             self.fail(f"Write API with default options raised an exception: {str(e)}")
+        finally:
+            client._write_api._on_complete()  # abort batch writes - otherwise test cycles through urllib3 retries
 
     def test_default_client(self):
         expected_precision = DefaultWriteOptions.write_precision.value
         expected_write_type = DefaultWriteOptions.write_type.value
 
+        import os
+        try:
+            os.environ["INFLUX_HOST"]
+        except KeyError:
+            os.environ["INFLUX_HOST"] = "http://my-influx.io"
+
+        try:
+            os.environ["INFLUX_TOKEN"]
+        except KeyError:
+            os.environ["INFLUX_TOKEN"] = "my-token"
+
+        try:
+            os.environ["INFLUX_DATABASE"]
+        except KeyError:
+            os.environ["INFLUX_DATABASE"] = "my-bucket"
+
         def verify_client_write_options(c):
             write_options = c._write_client_options.get('write_options')
             self.assertEqual(write_options.write_precision, expected_precision)
@@ -142,6 +161,21 @@ 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(InfluxDB3ClientQueryError) 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(InfluxDB3ClientQueryError) as err:
+                await c.query_async("SELECT * FROM my_data")
+            self.assertIn("Error while executing query", str(err.exception))
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_influxdb_client_3_integration.py b/tests/test_influxdb_client_3_integration.py
@@ -7,9 +7,9 @@
 
 import pyarrow
 import pytest
-from pyarrow._flight import FlightError
+from influxdb_client_3.exceptions import InfluxDB3ClientQueryError, InfluxDBError
 
-from influxdb_client_3 import InfluxDBClient3, InfluxDBError, write_client_options, WriteOptions
+from influxdb_client_3 import InfluxDBClient3, write_client_options, WriteOptions
 from tests.util import asyncio_run, lp_to_py_object
 
 
@@ -225,7 +225,7 @@ def test_queries_w_bad_cert(self):
                 query = "SELECT table_name FROM information_schema.tables"
                 client.query(query, mode="")
                 assert False, "query should throw SSL_ERROR"
-            except FlightError as fe:
+            except InfluxDB3ClientQueryError as fe:
                 assert str(fe).__contains__('SSL_ERROR_SSL')
             finally:
                 self.remove_test_cert(cert_file)
diff --git a/tests/util/mocks.py b/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,8 @@ 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

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# flake8: noqa`
	`2`	`+`
	`3`	`+from .exceptions import InfluxDB3ClientQueryError, InfluxDBError, InfluxDB3ClientError`