only add trace comment to sql if we are tracing

prisma · Sep 14, 2022 · b554aec · b554aec
1 parent 80833cf
commit b554aec
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git 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
@@ -33,7 +33,12 @@ macro_rules! sql_trace {
             // Temporary method to pass the traceid in an operation
             fn add_trace_id(self, trace_id: Option<String>) -> Self {
                 if let Some(traceparent) = trace_id {
-                    self.comment(format!("traceparent={}", traceparent))
+                    let out = should_sample(&traceparent);
+                    if should_sample(&traceparent) {
+                        self.comment(format!("traceparent={}", traceparent))
+                    } else {
+                        self
+                    }
                 } else {
                     self
                 }
@@ -42,6 +47,16 @@ macro_rules! sql_trace {
     };
 }
 
+fn should_sample(traceparent: &str) -> bool {
+    let trace_info: Vec<&str> = traceparent.split("-").collect();
+
+    if trace_info.len() != 4 {
+        false
+    } else {
+        trace_info.last() == Some(&"01")
+    }
+}
+
 sql_trace!(Insert<'_>);
 
 sql_trace!(Update<'_>);