From 34f162ec8569aaddd3c6dd53d69969e6a85ac1ae Mon Sep 17 00:00:00 2001 From: jacek-prisma Date: Tue, 10 Dec 2024 12:04:06 +0000 Subject: [PATCH] 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 --- Cargo.lock | 28 +++++++++++++++++++++++++--- quaint/Cargo.toml | 1 + quaint/src/connector/metrics.rs | 32 +++++++++++++++++++++++++++++--- quaint/src/tests/query.rs | 27 +++++++++++++++++++++++++++ 4 files changed, 82 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b31f4e1bd97..f26c8247b84 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "addr2line" @@ -2512,7 +2512,7 @@ dependencies = [ [[package]] name = "mongodb" version = "3.0.0" -source = "git+https://github.com/prisma/mongo-rust-driver.git?branch=RUST-1994/happy-eyeballs#31e0356391a7871bec000ae35fe0edc35582449e" +source = "git+https://github.com/prisma/mongo-rust-driver.git?branch=RUST-1994%2Fhappy-eyeballs#31e0356391a7871bec000ae35fe0edc35582449e" dependencies = [ "async-trait", "base64 0.13.1", @@ -2569,7 +2569,7 @@ dependencies = [ [[package]] name = "mongodb-internal-macros" version = "3.0.0" -source = "git+https://github.com/prisma/mongo-rust-driver.git?branch=RUST-1994/happy-eyeballs#31e0356391a7871bec000ae35fe0edc35582449e" +source = "git+https://github.com/prisma/mongo-rust-driver.git?branch=RUST-1994%2Fhappy-eyeballs#31e0356391a7871bec000ae35fe0edc35582449e" dependencies = [ "proc-macro2", "quote", @@ -3742,6 +3742,7 @@ dependencies = [ "tokio-util 0.7.8", "tracing", "tracing-futures", + "tracing-test", "url", "uuid", "ws_stream_tungstenite", @@ -6038,6 +6039,27 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "tracing-test" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "557b891436fe0d5e0e363427fc7f217abf9ccd510d5136549847bdcbcd011d68" +dependencies = [ + "tracing-core", + "tracing-subscriber", + "tracing-test-macro", +] + +[[package]] +name = "tracing-test-macro" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04659ddb06c87d233c566112c1c9c5b9e98256d9af50ec3bc9c8327f873a7568" +dependencies = [ + "quote", + "syn 2.0.58", +] + [[package]] name = "try-lock" version = "0.2.4" diff --git a/quaint/Cargo.toml b/quaint/Cargo.toml index 4e13c764b22..7d066e74466 100644 --- a/quaint/Cargo.toml +++ b/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"] } diff --git a/quaint/src/connector/metrics.rs b/quaint/src/connector/metrics.rs index 78fc7f99c72..37143866a67 100644 --- a/quaint/src/connector/metrics.rs +++ b/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>, { - 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}") + } +} diff --git a/quaint/src/tests/query.rs b/quaint/src/tests/query.rs index 91fb5b2985a..d44cada88e6 100644 --- a/quaint/src/tests/query.rs +++ b/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(()) +}