Merge pull request #58 from InfluxCommunity/57-add-forced-trace-debug

Jayclifford345 · web-flow · commit 7e6bc49025af · 2023-09-15T16:03:24.000+01:00
added trace debugging feature for support
diff --git 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
@@ -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
