random id to traceId headers (jupyter-server#239)

Zsailer · Nov 9, 2021 · 41804d8 · 41804d8
1 parent 401529d
commit 41804d8
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/data_studio_jupyter_extensions/configurables/notebook_service.py b/data_studio_jupyter_extensions/configurables/notebook_service.py
@@ -1,5 +1,6 @@
 import base64
 import json
+import secrets
 import ssl
 import urllib
 
@@ -95,15 +96,20 @@ async def fetch_token(self):  # pragma: no cover
         token = resp_body["access_token"]
         return token
 
-    @property
-    def headers(self):
+    def get_headers(self, trace_id=None):
         """Get the headers needed for a request"""
-        return {
+        headers = {
             "Authorization": "Bearer %s" % self.request_token,
             "Content-Type": "application/json;charset=UTF-8",
             "datastudio-token-provider": "IAM-CORP",
-            "X-B3-TraceId": "b34532b65f9fd4ee",
         }
+        # Generate a TraceID if not given.
+        if trace_id:
+            headers["X-B3-TraceId"] = trace_id
+        else:
+            # Generate a random character string that's 32 characters long.
+            headers["X-B3-TraceId"] = secrets.token_hex(16)
+        return headers
 
     @property
     def base_url(self):
@@ -114,7 +120,7 @@ def _get_request(self, url, method, data):
         return HTTPRequest(
             url=url,
             method=method.upper(),
-            headers=self.headers,
+            headers=self.get_headers(),
             body=json.dumps(data),
             allow_nonstandard_methods=True,
             ca_certs=self.ssl_cert_file,