From b94773c732a3362a89027b552dc67648a365ef3b Mon Sep 17 00:00:00 2001 From: Jay Clifford Date: Fri, 15 Sep 2023 16:02:55 +0100 Subject: [PATCH] added trace debugging feature for support --- influxdb_client_3/__init__.py | 9 +++++++-- influxdb_client_3/debug.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) create mode 100644 influxdb_client_3/debug.py diff --git a/influxdb_client_3/__init__.py b/influxdb_client_3/__init__.py index 6b518bb..f8ce4d5 100644 --- a/influxdb_client_3/__init__.py +++ b/influxdb_client_3/__init__.py @@ -8,6 +8,7 @@ from influxdb_client_3.read_file import UploadFile + def write_client_options(**kwargs): """ Function for providing additional arguments for the WriteApi client. @@ -158,7 +159,7 @@ def _process_dataframe(self, df, measurement_name, tag_columns, timestamp_column data_frame_tag_columns=tag_columns, data_frame_timestamp_column=timestamp_column, **kwargs) - def query(self, query, language="sql", mode="all", database=None, **kwargs ): + def query(self, query, language="sql", mode="all", database=None,**kwargs ): """ Query data from InfluxDB. @@ -173,12 +174,16 @@ def query(self, query, language="sql", mode="all", database=None, **kwargs ): :param kwargs: Additional arguments for the query. :return: The queried data. """ + + if database is None: database = self._database try: + headers = [(b"authorization", f"Bearer {self._token}".encode('utf-8'))] + # Create an authorization header - _options = FlightCallOptions(headers=[(b"authorization", f"Bearer {self._token}".encode('utf-8'))], **kwargs) + _options = FlightCallOptions(headers=headers, **kwargs) ticket_data = {"database": database, "sql_query": query, "query_type": language} ticket = Ticket(json.dumps(ticket_data).encode('utf-8')) flight_reader = self._flight_client.do_get(ticket, _options) diff --git a/influxdb_client_3/debug.py b/influxdb_client_3/debug.py new file mode 100644 index 0000000..73cd866 --- /dev/null +++ b/influxdb_client_3/debug.py @@ -0,0 +1,19 @@ +import pyarrow.flight as flight + +class TracingClientMiddleWareFactory(flight.ClientMiddleware): + def start_call(self, info): + print("Starting new call:", info) + return TracingClientMiddleware() + +class TracingClientMiddleware(flight.ClientMiddleware): + def sending_headers(self): + print("Sending trace ID:", "traceheader") + return { + "x-tracing-id": "traceheader", + } + + def received_headers(self, headers): + if "trace-id" in headers: + trace_id = headers["trace-id"][0] + print("Found trace header with value:", trace_id) + # Don't overwrite our trace ID \ No newline at end of file