diff --git a/query-engine/connectors/sql-query-connector/src/query_ext.rs b/query-engine/connectors/sql-query-connector/src/query_ext.rs index 2ed1ce32b583..3d52e22af2e7 100644 --- a/query-engine/connectors/sql-query-connector/src/query_ext.rs +++ b/query-engine/connectors/sql-query-connector/src/query_ext.rs @@ -50,9 +50,7 @@ pub trait QueryExt: Queryable + Send + Sync { Query::Select(Box::from(x.comment(trace_parent_to_string(span_ctx)))) } // This is part of the required changes to pass a traceid - (Query::Select(x), Some(traceparent)) => { - Query::Select(Box::from(x.comment(format!("traceparent={}", traceparent)))) - } + (Query::Select(x), trace_id) => Query::Select(Box::from(x.add_trace_id(trace_id))), (q, _) => q, }; diff --git a/query-engine/connectors/sql-query-connector/src/sql_trace.rs b/query-engine/connectors/sql-query-connector/src/sql_trace.rs index c3d67b46fc9e..29cab8198215 100644 --- a/query-engine/connectors/sql-query-connector/src/sql_trace.rs +++ b/query-engine/connectors/sql-query-connector/src/sql_trace.rs @@ -33,7 +33,11 @@ macro_rules! sql_trace { // Temporary method to pass the traceid in an operation fn add_trace_id(self, trace_id: Option) -> Self { if let Some(traceparent) = trace_id { - self.comment(format!("traceparent={}", traceparent)) + if should_sample(&traceparent) { + self.comment(format!("traceparent={}", traceparent)) + } else { + self + } } else { self } @@ -42,6 +46,10 @@ macro_rules! sql_trace { }; } +fn should_sample(traceparent: &str) -> bool { + traceparent.split('-').count() == 4 && traceparent.ends_with("-01") +} + sql_trace!(Insert<'_>); sql_trace!(Update<'_>);