feat: update user-agent and improve transport selection in queryClient

.gitignore

@@ -9,3 +9,7 @@ pyinflux3*.egg-info
 __pycache__
 .idea
 *.egg-info/
+temp/
+test-reports/
+coverage.xml
+.coverage

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.6.0 [unreleased]
 
+### Features
+
+1. [#92](https://github.com/InfluxCommunity/influxdb3-python/pull/92): Update `user-agent` header value to `influxdb3-python/{VERSION}` and add it to queries as well. 
+
 ## 0.5.0 [2024-05-17]
 
 ### Features

influxdb_client_3/__init__.py

@@ -10,6 +10,7 @@
 from influxdb_client_3.write_client.client.write_api import WriteApi as _WriteApi, SYNCHRONOUS, ASYNCHRONOUS, \
     PointSettings
 from influxdb_client_3.write_client.domain.write_precision import WritePrecision
+from influxdb_client_3.version import USER_AGENT
 
 try:
     import polars as pl
@@ -147,7 +148,19 @@ def __init__(
 
         if query_port_overwrite is not None:
             port = query_port_overwrite
-        self._flight_client = FlightClient(f"grpc+tls://{hostname}:{port}", **self._flight_client_options)
+
+        gen_opts = [
+            ("grpc.secondary_user_agent", USER_AGENT)
+        ]
+
+        self._flight_client_options["generic_options"] = gen_opts
+
+        if scheme == 'https':
+            connection_string = f"grpc+tls://{hostname}:{port}"
+        else:
+            connection_string = f"grpc+tcp://{hostname}:{port}"
+
+        self._flight_client = FlightClient(connection_string, **self._flight_client_options)
 
     def write(self, record=None, database=None, **kwargs):
         """

influxdb_client_3/version.py

@@ -0,0 +1,4 @@
+"""Version of the Client that is used in User-Agent header."""
+
+VERSION = '0.6.0dev0'
+USER_AGENT = f'influxdb3-python/{VERSION}'

influxdb_client_3/write_client/__init__.py

@@ -27,5 +27,5 @@
 from influxdb_client_3.write_client.domain.write_precision import WritePrecision
 
 from influxdb_client_3.write_client.configuration import Configuration
-from influxdb_client_3.write_client.version import VERSION
+from influxdb_client_3.version import VERSION
 __version__ = VERSION

influxdb_client_3/write_client/_sync/api_client.py

@@ -76,8 +76,8 @@ def __init__(self, configuration=None, header_name=None, header_value=None,
             self.default_headers[header_name] = header_value
         self.cookie = cookie
         # Set default User-Agent.
-        from influxdb_client_3.write_client.version import VERSION
-        self.user_agent = f'influxdb-client-python/{VERSION}'
+        from influxdb_client_3.version import USER_AGENT
+        self.user_agent = USER_AGENT
 
     def __del__(self):
         """Dispose pools."""

tests/__init__.py

@@ -0,0 +1 @@
+# needed to resolve some module imports when running pytest

tests/test_api_client.py

@@ -0,0 +1,71 @@
+import unittest
+from unittest import mock
+
+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.service import WriteService
+from influxdb_client_3.version import VERSION
+
+
+_package = "influxdb3-python"
+_sentHeaders = {}
+
+
+def mock_rest_request(method,
+                      url,
+                      query_params=None,
+                      headers=None,
+                      body=None,
+                      post_params=None,
+                      _preload_content=True,
+                      _request_timeout=None,
+                      **urlopen_kw):
+    class MockResponse:
+        def __init__(self, data, status_code):
+            self.data = data
+            self.status_code = status_code
+
+        def data(self):
+            return self.data
+
+    global _sentHeaders
+    _sentHeaders = headers
+
+    return MockResponse(None, 200)
+
+
+class ApiClientTests(unittest.TestCase):
+
+    def test_default_headers(self):
+        global _package
+        conf = Configuration()
+        client = ApiClient(conf,
+                           header_name="Authorization",
+                           header_value="Bearer TEST_TOKEN")
+        self.assertIsNotNone(client.default_headers["User-Agent"])
+        self.assertIsNotNone(client.default_headers["Authorization"])
+        self.assertEqual(f"{_package}/{VERSION}", client.default_headers["User-Agent"])
+        self.assertEqual("Bearer TEST_TOKEN", client.default_headers["Authorization"])
+
+    @mock.patch("influxdb_client_3.write_client._sync.rest.RESTClientObject.request",
+                side_effect=mock_rest_request)
+    def test_call_api(self, mock_post):
+        global _package
+        global _sentHeaders
+        _sentHeaders = {}
+
+        conf = Configuration()
+        client = ApiClient(conf,
+                           header_name="Authorization",
+                           header_value="Bearer TEST_TOKEN")
+        service = WriteService(client)
+        service.post_write("TEST_ORG", "TEST_BUCKET", "data,foo=bar val=3.14")
+        self.assertEqual(4, len(_sentHeaders.keys()))
+        self.assertIsNotNone(_sentHeaders["Accept"])
+        self.assertEqual("application/json", _sentHeaders["Accept"])
+        self.assertIsNotNone(_sentHeaders["Content-Type"])
+        self.assertEqual("text/plain", _sentHeaders["Content-Type"])
+        self.assertIsNotNone(_sentHeaders["Authorization"])
+        self.assertEqual("Bearer TEST_TOKEN", _sentHeaders["Authorization"])
+        self.assertIsNotNone(_sentHeaders["User-Agent"])
+        self.assertEqual(f"{_package}/{VERSION}", _sentHeaders["User-Agent"])

tests/test_query.py

@@ -1,9 +1,115 @@
 import unittest
+import struct
 from unittest.mock import Mock, patch, ANY
 
-from pyarrow.flight import Ticket
+from pyarrow import (
+    array,
+    Table
+)
+
+from pyarrow.flight import (
+    FlightServerBase,
+    FlightUnauthenticatedError,
+    GeneratorStream,
+    ServerMiddleware,
+    ServerMiddlewareFactory,
+    ServerAuthHandler,
+    Ticket
+)
 
 from influxdb_client_3 import InfluxDBClient3
+from influxdb_client_3.version import USER_AGENT
+
+
+def case_insensitive_header_lookup(headers, lkey):
+    """Lookup the value of a given key in the given headers.
+       The lkey is case-insensitive.
+    """
+    for key in headers:
+        if key.lower() == lkey.lower():
+            return headers.get(key)
+
+
+class NoopAuthHandler(ServerAuthHandler):
+    """A no-op auth handler - as seen in pyarrow tests"""
+
+    def authenticate(self, outgoing, incoming):
+        """Do nothing"""
+
+    def is_valid(self, token):
+        """
+        Return an empty string
+        N.B. Returning None causes Type error
+        :param token:
+        :return:
+        """
+        return ""
+
+
+_req_headers = {}
+
+
+class HeaderCheckServerMiddlewareFactory(ServerMiddlewareFactory):
+    """Factory to create HeaderCheckServerMiddleware and check header values"""
+    def start_call(self, info, headers):
+        auth_header = case_insensitive_header_lookup(headers, "Authorization")
+        values = auth_header[0].split(' ')
+        if values[0] != 'Bearer':
+            raise FlightUnauthenticatedError("Token required")
+        global _req_headers
+        _req_headers = headers
+        return HeaderCheckServerMiddleware(values[1])
+
+
+class HeaderCheckServerMiddleware(ServerMiddleware):
+    """
+    Middleware needed to catch request headers via factory
+    N.B. As found in pyarrow tests
+    """
+    def __init__(self, token):
+        self.token = token
+
+    def sending_headers(self):
+        return {'authorization': 'Bearer ' + self.token}
+
+
+class HeaderCheckFlightServer(FlightServerBase):
+    """Mock server handle gRPC do_get calls"""
+    def do_get(self, context, ticket):
+        """Return something to avoid needless errors"""
+        data = [
+            array([b"Vltava", struct.pack('<i', 105), b"FM"])
+        ]
+        table = Table.from_arrays(data, names=['a'])
+        return GeneratorStream(
+            table.schema,
+            self.number_batches(table),
+            options={})
+
+    @staticmethod
+    def number_batches(table):
+        for idx, batch in enumerate(table.to_batches()):
+            buf = struct.pack('<i', idx)
+            yield batch, buf
+
+
+def test_influx_default_query_headers():
+    with HeaderCheckFlightServer(
+            auth_handler=NoopAuthHandler(),
+            middleware={"check": HeaderCheckServerMiddlewareFactory()}) as server:
+        global _req_headers
+        _req_headers = {}
+        client = InfluxDBClient3(
+            host=f'http://localhost:{server.port}',
+            org='test_org',
+            databse='test_db',
+            token='TEST_TOKEN'
+        )
+        client.query('SELECT * FROM test')
+        assert len(_req_headers) > 0
+        assert _req_headers['authorization'][0] == "Bearer TEST_TOKEN"
+        assert _req_headers['user-agent'][0].find(USER_AGENT) > -1
+        _req_headers = {}
 
 
 class QueryTests(unittest.TestCase):