diff --git a/dagshub/common/config.py b/dagshub/common/config.py index 3b932c6f..e82b0063 100644 --- a/dagshub/common/config.py +++ b/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" diff --git a/dagshub/common/tracing.py b/dagshub/common/tracing.py new file mode 100644 index 00000000..97e6579d --- /dev/null +++ b/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}" diff --git a/dagshub/data_engine/client/data_client.py b/dagshub/data_engine/client/data_client.py index 5ae8eb25..d5c61e2f 100644 --- a/dagshub/data_engine/client/data_client.py +++ b/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(