fix: strip traceparent comment from info_span (#5078)

* fix: strip traceparent comment from info_span

* fix: handle non-eoi traceparent coment

* fix: update assertion check

* fix: assert on logs_contain

* fix: simplify test to avoid db differences

* fix: fix a lint

quaint/Cargo.toml

@@ -110,6 +110,7 @@ quaint-test-macros = { path = "quaint-test-macros" }
 quaint-test-setup = { path = "quaint-test-setup" }
 tokio = { version = "1", features = ["macros", "time"] }
 expect-test = "1"
+tracing-test = "0.2"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 getrandom = { workspace = true, features = ["js"] }

quaint/src/connector/metrics.rs

@@ -3,7 +3,7 @@ use tracing::{info_span, Instrument};
 
 use crate::ast::{Params, Value};
 use crosstarget_utils::time::ElapsedTimeCounter;
-use std::future::Future;
+use std::{fmt, future::Future};
 
 pub async fn query<'a, F, T, U>(
     tag: &'static str,
@@ -16,8 +16,12 @@ where
     F: FnOnce() -> U + 'a,
     U: Future<Output = crate::Result<T>>,
 {
-    let span =
-        info_span!("quaint:query", "db.system" = db_system_name, "db.statement" = %query, "otel.kind" = "client");
+    let span = info_span!(
+        "quaint:query",
+        "db.system" = db_system_name,
+        "db.statement" = %QueryForTracing(query),
+        "otel.kind" = "client"
+    );
     do_query(tag, query, params, f).instrument(span).await
 }
 
@@ -97,3 +101,25 @@ fn trace_query<'a>(query: &'a str, params: &'a [Value<'_>], result: &str, start:
         duration_ms = start.elapsed_time().as_millis() as u64,
     );
 }
+
+struct QueryForTracing<'a>(&'a str);
+
+impl fmt::Display for QueryForTracing<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let query = self
+            .0
+            .split_once("/* traceparent=")
+            .map_or(self.0, |(str, remainder)| {
+                if remainder
+                    .split_once("*/")
+                    .is_some_and(|(_, suffix)| suffix.trim_end().is_empty())
+                {
+                    str
+                } else {
+                    self.0
+                }
+            })
+            .trim();
+        write!(f, "{query}")
+    }
+}

quaint/src/tests/query.rs

@@ -10,6 +10,7 @@ use crate::{
 };
 use quaint_test_macros::test_each_connector;
 use quaint_test_setup::Tags;
+use tracing_test::traced_test;
 
 #[test_each_connector]
 async fn single_value(api: &mut dyn TestApi) -> crate::Result<()> {
@@ -3636,3 +3637,29 @@ async fn overflowing_int_errors_out(api: &mut dyn TestApi) -> crate::Result<()>
 
     Ok(())
 }
+
+#[test_each_connector]
+#[traced_test]
+async fn traceparent_is_stripped_from_the_log(api: &mut dyn TestApi) -> crate::Result<()> {
+    api.conn()
+        .query_raw("SELECT 1 /* traceparent=1 */", &[])
+        .await?
+        .into_single()?;
+    let expected = r#"db.statement=SELECT 1 otel.kind="client""#.to_owned();
+    assert!(logs_contain(&expected), "expected logs to contain '{expected}'");
+
+    Ok(())
+}
+
+#[test_each_connector]
+#[traced_test]
+async fn traceparent_inside_of_query_isnt_stripped_from_log(api: &mut dyn TestApi) -> crate::Result<()> {
+    api.conn()
+        .query_raw("SELECT /* traceparent=1 */ 1", &[])
+        .await?
+        .into_single()?;
+    let expected = r#"db.statement=SELECT /* traceparent=1 */ 1 otel.kind="client""#.to_owned();
+    assert!(logs_contain(&expected), "expected logs to contain '{expected}'");
+
+    Ok(())
+}