Add support for W3C Trace Context in graphql requests

dagshub/common/config.py

@@ -20,6 +20,7 @@
 DAGSHUB_PASSWORD_KEY = "DAGSHUB_PASSWORD"
 HTTP_TIMEOUT_KEY = "DAGSHUB_HTTP_TIMEOUT"
 DAGSHUB_QUIET_KEY = "DAGSHUB_QUIET"
+DISABLE_TRACEPARENT_KEY = "DAGSHUB_DISABLE_TRACEPARENT"
 
 
 def set_host(new_host: str):
@@ -48,6 +49,8 @@ def set_host(new_host: str):
 
 quiet = bool(os.environ.get(DAGSHUB_QUIET_KEY, False))
 
+disable_traceparent = bool(os.environ.get(DISABLE_TRACEPARENT_KEY, False))
+
 # DVC config templates
 CONFIG_GITIGNORE = "/config.local\n/tmp\n/cache"
 

dagshub/common/tracing.py

@@ -0,0 +1,26 @@
+import secrets
+
+
+def _generate_non_zero_hex(byte_count: int) -> str:
+    """
+    Generate a non-zero hex string of the given byte length.
+
+    Per W3C Trace Context, trace-id and parent-id MUST NOT be all zeros.
+    """
+    while True:
+        value = secrets.token_hex(byte_count)
+        if int(value, 16) != 0:
+            return value
+
+
+def build_traceparent() -> str:
+    """
+    Build a W3C Trace Context traceparent header value.
+
+    Format: version(2)-trace-id(32)-parent-id(16)-flags(2)
+    """
+    version = "00"
+    trace_id = _generate_non_zero_hex(16)
+    parent_id = _generate_non_zero_hex(8)
+    flags = "01"  # sampled
+    return f"{version}-{trace_id}-{parent_id}-{flags}"

dagshub/data_engine/client/data_client.py

@@ -5,12 +5,13 @@
 import dacite
 import gql
 import rich.progress
-from gql.transport.exceptions import TransportQueryError
+from gql.transport.exceptions import TransportQueryError, TransportServerError
 from gql.transport.requests import RequestsHTTPTransport
 
 import dagshub.auth
 import dagshub.common.config
 from dagshub.common import config
+from dagshub.common.tracing import build_traceparent
 from dagshub.common.analytics import send_analytics_event
 from dagshub.common.rich_util import get_rich_progress
 from dagshub.data_engine.client.gql_introspections import GqlIntrospections, TypesIntrospection
@@ -184,10 +185,18 @@ def _exec(
         if validate:
             query.validate_params(params if params else {}, self.query_introspection)
         q = gql.gql(query.generate())
+        headers = dict(config.requests_headers)
+        traceparent = None
+        if not config.disable_traceparent:
+            traceparent = build_traceparent()
+            headers["traceparent"] = traceparent
         try:
-            resp = self.client.execute(q, variable_values=params)
-        except TransportQueryError as e:
-            raise DataEngineGqlError(e, self.client.transport.response_headers.get("X-DagsHub-Support-Id"))
+            resp = self.client.execute(q, variable_values=params, extra_args={'headers': headers})
+        except (TransportQueryError, TransportServerError) as e:
+            support_id = self.client.transport.response_headers.get("X-DagsHub-Support-Id")
+            if support_id is None:
+                support_id = traceparent
+            raise DataEngineGqlError(e, support_id)
         return resp
 
     def _datasource_query(