From 8aaf982d17922ddbbb02b7179090c22b81bc4b22 Mon Sep 17 00:00:00 2001 From: karel rehor Date: Thu, 6 Jun 2024 11:07:35 +0200 Subject: [PATCH 1/4] feat: updates user-agent header, improves query transport scheme selection --- .gitignore | 4 + influxdb_client_3/__init__.py | 13 +- .../{write_client => }/version.py | 3 +- influxdb_client_3/write_client/__init__.py | 2 +- .../write_client/_sync/api_client.py | 4 +- tests/__init__.py | 1 + tests/test_api_client.py | 71 +++++++++++ tests/test_query.py | 111 +++++++++++++++++- 8 files changed, 203 insertions(+), 6 deletions(-) rename influxdb_client_3/{write_client => }/version.py (50%) create mode 100644 tests/__init__.py create mode 100644 tests/test_api_client.py diff --git a/.gitignore b/.gitignore index 85224cc..c6c28d1 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,7 @@ pyinflux3*.egg-info __pycache__ .idea *.egg-info/ +temp/ +test-reports/ +coverage.xml +.coverage diff --git a/influxdb_client_3/__init__.py b/influxdb_client_3/__init__.py index 7bb877b..4c34def 100644 --- a/influxdb_client_3/__init__.py +++ b/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, VERSION try: import polars as pl @@ -147,7 +148,17 @@ 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': + self._flight_client = FlightClient(f"grpc+tls://{hostname}:{port}", **self._flight_client_options) + else: + self._flight_client = FlightClient(f"grpc+tcp://{hostname}:{port}", **self._flight_client_options) def write(self, record=None, database=None, **kwargs): """ diff --git a/influxdb_client_3/write_client/version.py b/influxdb_client_3/version.py similarity index 50% rename from influxdb_client_3/write_client/version.py rename to influxdb_client_3/version.py index b8eb2f2..b875dd1 100644 --- a/influxdb_client_3/write_client/version.py +++ b/influxdb_client_3/version.py @@ -1,3 +1,4 @@ """Version of the Client that is used in User-Agent header.""" -VERSION = '1.38.0dev0' +VERSION = '0.6.0dev0' +USER_AGENT = f'influxdb3-python/{VERSION}' \ No newline at end of file diff --git a/influxdb_client_3/write_client/__init__.py b/influxdb_client_3/write_client/__init__.py index efaf933..8abb935 100644 --- a/influxdb_client_3/write_client/__init__.py +++ b/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 diff --git a/influxdb_client_3/write_client/_sync/api_client.py b/influxdb_client_3/write_client/_sync/api_client.py index cc22ff8..b72a840 100644 --- a/influxdb_client_3/write_client/_sync/api_client.py +++ b/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.""" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..035cccd --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# needed to resolve some module imports when running pytest diff --git a/tests/test_api_client.py b/tests/test_api_client.py new file mode 100644 index 0000000..edffe42 --- /dev/null +++ b/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"]) diff --git a/tests/test_query.py b/tests/test_query.py index e2131f6..3225455 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -1,9 +1,118 @@ 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 ( + FlightCallOptions, + FlightClient, + 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(' 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): From 5f693331e73f46fdb2a5432cca5c99f09edce963 Mon Sep 17 00:00:00 2001 From: karel rehor Date: Thu, 6 Jun 2024 13:11:06 +0200 Subject: [PATCH 2/4] chore: fix lint issues --- influxdb_client_3/__init__.py | 4 ++-- influxdb_client_3/version.py | 2 +- tests/test_query.py | 7 ++----- 3 files changed, 5 insertions(+), 8 deletions(-) diff --git a/influxdb_client_3/__init__.py b/influxdb_client_3/__init__.py index 4c34def..a4407ee 100644 --- a/influxdb_client_3/__init__.py +++ b/influxdb_client_3/__init__.py @@ -10,7 +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, VERSION +from influxdb_client_3.version import USER_AGENT try: import polars as pl @@ -150,7 +150,7 @@ def __init__( port = query_port_overwrite gen_opts = [ - ("grpc.secondary_user_agent",USER_AGENT) + ("grpc.secondary_user_agent", USER_AGENT) ] self._flight_client_options["generic_options"] = gen_opts diff --git a/influxdb_client_3/version.py b/influxdb_client_3/version.py index b875dd1..e7a51c4 100644 --- a/influxdb_client_3/version.py +++ b/influxdb_client_3/version.py @@ -1,4 +1,4 @@ """Version of the Client that is used in User-Agent header.""" VERSION = '0.6.0dev0' -USER_AGENT = f'influxdb3-python/{VERSION}' \ No newline at end of file +USER_AGENT = f'influxdb3-python/{VERSION}' diff --git a/tests/test_query.py b/tests/test_query.py index 3225455..1fd7cb8 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -8,8 +8,6 @@ ) from pyarrow.flight import ( - FlightCallOptions, - FlightClient, FlightServerBase, FlightUnauthenticatedError, GeneratorStream, @@ -80,7 +78,7 @@ class HeaderCheckFlightServer(FlightServerBase): def do_get(self, context, ticket): """Return something to avoid needless errors""" data = [ - array([b"Vltava", struct.pack(' Date: Thu, 6 Jun 2024 13:47:11 +0200 Subject: [PATCH 3/4] docs: updates CHANGELOG.md --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2148310..1e7aa49 100644 --- a/CHANGELOG.md +++ b/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 From 8e3197cd8d3937ded17597b06f19b29c6d53dc07 Mon Sep 17 00:00:00 2001 From: karel rehor Date: Thu, 6 Jun 2024 14:39:28 +0200 Subject: [PATCH 4/4] chore: improve _flight_client init with connection_string --- influxdb_client_3/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/influxdb_client_3/__init__.py b/influxdb_client_3/__init__.py index a4407ee..b01d0e1 100644 --- a/influxdb_client_3/__init__.py +++ b/influxdb_client_3/__init__.py @@ -156,9 +156,11 @@ def __init__( self._flight_client_options["generic_options"] = gen_opts if scheme == 'https': - self._flight_client = FlightClient(f"grpc+tls://{hostname}:{port}", **self._flight_client_options) + connection_string = f"grpc+tls://{hostname}:{port}" else: - self._flight_client = FlightClient(f"grpc+tcp://{hostname}:{port}", **self._flight_client_options) + 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): """