es_entity/
sql_commenter.rs

1//! Correlate SQL statements with OpenTelemetry traces.
2//!
3//! [`annotate_sql`] appends a [sqlcommenter](https://google.github.io/sqlcommenter)-style
4//! `traceparent` comment carrying the currently active span's trace context to
5//! a SQL statement, e.g.
6//!
7//! ```text
8//! SELECT * FROM users WHERE id = $1 /*traceparent='00-<32-hex>-<16-hex>-01'*/
9//! ```
10//!
11//! The comment survives Postgres query normalization and appears in
12//! `pg_stat_activity` and in Postgres logs (slow query log, `auto_explain`,
13//! lock waits), so a statement observed server-side can be matched back to a
14//! distributed trace. `pg_stat_statements` retains the comment of the *first*
15//! execution seen per query id (comments are excluded from the query id
16//! jumble), which yields an exemplar trace per statement shape.
17//!
18//! # Sampling and the prepared statement cache
19//!
20//! Annotation only happens when the active span is **sampled**: an un-sampled
21//! span is never exported, so its trace id cannot be looked up in the tracing
22//! backend and annotating would be pure cost. The cost matters because an
23//! annotated statement's text is unique per span — it can never be served
24//! from sqlx's per-connection prepared statement cache and is executed
25//! non-persistently (a server-side parse + plan per execution). Gating on the
26//! sampled flag keeps full prepared-statement reuse for all un-sampled
27//! traffic.
28//!
29//! Annotation requires the `tracing-context` feature and an initialized
30//! OpenTelemetry tracing layer. Otherwise [`annotate_sql`] is a zero-cost
31//! pass-through.
32
33use std::borrow::Cow;
34
35/// The W3C `trace-flags` sampled bit.
36#[cfg(feature = "tracing-context")]
37const SAMPLED: u8 = 0x01;
38
39/// W3C `traceparent` of the currently active span, e.g.
40/// `00-<32-hex-trace-id>-<16-hex-span-id>-01`, if that span is sampled.
41///
42/// Returns `None` when there is no valid span context (outside any span, or
43/// tracing/OTEL not initialized), when the span is not sampled (its trace is
44/// never exported, so the annotation could not be correlated with anything),
45/// or when the `tracing-context` feature is disabled.
46#[cfg(feature = "tracing-context")]
47pub fn current_traceparent() -> Option<String> {
48    crate::context::TracingContext::current()
49        .filter(|ctx| ctx.trace_flags & SAMPLED != 0)
50        .map(|ctx| ctx.traceparent)
51}
52
53/// W3C `traceparent` of the currently active span. Always `None` without the
54/// `tracing-context` feature.
55#[cfg(not(feature = "tracing-context"))]
56pub fn current_traceparent() -> Option<String> {
57    None
58}
59
60/// Appends the current sampled span's `traceparent` as a SQL comment to `sql`.
61///
62/// Returns [`Cow::Borrowed`] (no allocation) when there is no sampled span
63/// context; callers can use this to skip any annotation-dependent work.
64pub fn annotate_sql(sql: &str) -> Cow<'_, str> {
65    match current_traceparent() {
66        Some(traceparent) => Cow::Owned(format!("{sql} /*traceparent='{traceparent}'*/")),
67        None => Cow::Borrowed(sql),
68    }
69}
70
71#[cfg(all(test, feature = "tracing-context"))]
72mod tests {
73    use opentelemetry::trace::TracerProvider as _;
74    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
75
76    use super::*;
77
78    fn subscriber_with_sampler(
79        sampler: opentelemetry_sdk::trace::Sampler,
80    ) -> tracing::subscriber::DefaultGuard {
81        let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
82            .with_sampler(sampler)
83            .build();
84        let tracer = provider.tracer("sql-commenter-test");
85        tracing_subscriber::registry()
86            .with(tracing_opentelemetry::layer().with_tracer(tracer))
87            .set_default()
88    }
89
90    #[test]
91    fn no_span_context_is_noop() {
92        assert!(current_traceparent().is_none());
93        assert_eq!(annotate_sql("SELECT 1"), "SELECT 1");
94    }
95
96    #[test]
97    fn annotates_within_sampled_span() {
98        let _subscriber = subscriber_with_sampler(opentelemetry_sdk::trace::Sampler::AlwaysOn);
99
100        let span = tracing::info_span!("test_span");
101        let _guard = span.enter();
102
103        let tp = current_traceparent().expect("should have traceparent in span");
104        let parts: Vec<&str> = tp.split('-').collect();
105        assert_eq!(parts.len(), 4);
106        assert_eq!(parts[0], "00");
107        assert_eq!(parts[1].len(), 32);
108        assert_eq!(parts[2].len(), 16);
109        assert_eq!(parts[3], "01");
110
111        assert_eq!(
112            annotate_sql("SELECT 1"),
113            format!("SELECT 1 /*traceparent='{tp}'*/")
114        );
115    }
116
117    #[test]
118    fn unsampled_span_is_not_annotated() {
119        let _subscriber = subscriber_with_sampler(opentelemetry_sdk::trace::Sampler::AlwaysOff);
120
121        let span = tracing::info_span!("test_span");
122        let _guard = span.enter();
123
124        assert!(current_traceparent().is_none());
125        assert_eq!(annotate_sql("SELECT 1"), "SELECT 1");
126    }
127}
es_entity/sql_commenter.rs

es_entity/
sql_commenter.rs
