sqlx_otel/
annotations.rs

1use crate::pool::SharedState;
2
3/// Generate `with_annotations` and `with_operation` inherent methods for wrapper types
4/// that hold a `SharedState` and return `AnnotatedMut` via mutable borrows.
5///
6/// Invoke inside an `impl` block – the macro emits only method items, not the enclosing
7/// `impl`.
8macro_rules! impl_with_annotations_mut {
9    () => {
10        /// Return an annotated executor that attaches per-query semantic convention
11        /// attributes to every span created by the next operation.
12        ///
13        /// The returned wrapper borrows `self` mutably and implements `sqlx::Executor`
14        /// with the same instrumentation, but with annotation values threaded through to
15        /// span creation.
16        #[must_use]
17        pub fn with_annotations(
18            &mut self,
19            annotations: crate::annotations::QueryAnnotations,
20        ) -> crate::annotations::AnnotatedMut<'_, Self> {
21            crate::annotations::AnnotatedMut {
22                state: self.state.clone(),
23                annotations,
24                inner: self,
25            }
26        }
27
28        /// Shorthand for annotating the next operation with `db.operation.name` and
29        /// `db.collection.name`.
30        ///
31        /// Equivalent to
32        /// `self.with_annotations(QueryAnnotations::new().operation(op).collection(coll))`.
33        #[must_use]
34        pub fn with_operation(
35            &mut self,
36            operation: impl Into<String>,
37            collection: impl Into<String>,
38        ) -> crate::annotations::AnnotatedMut<'_, Self> {
39            self.with_annotations(
40                crate::annotations::QueryAnnotations::new()
41                    .operation(operation)
42                    .collection(collection),
43            )
44        }
45    };
46}
47
48/// Per-query annotation values that enrich OpenTelemetry spans with semantic-convention
49/// attributes the library cannot derive automatically (because it does not parse SQL).
50///
51/// Use the builder methods to set whichever attributes apply to a given query, then pass
52/// the result through one of the equivalent annotation surfaces:
53///
54/// - **Executor-side** – [`Pool::with_annotations`](crate::Pool::with_annotations),
55///   [`PoolConnection::with_annotations`](crate::PoolConnection::with_annotations), or
56///   [`Transaction::with_annotations`](crate::Transaction::with_annotations). Returns a
57///   borrowed wrapper that is itself an `sqlx::Executor`.
58/// - **Query-side** – [`QueryAnnotateExt::with_annotations`](crate::QueryAnnotateExt) on
59///   the builder produced by `sqlx::query`, `sqlx::query_as`, `sqlx::query_scalar`, or
60///   their `_!` macro forms.
61///
62/// Both surfaces produce identical telemetry; pick whichever keeps the annotation closer to
63/// the thing it describes.
64///
65/// # Example
66///
67/// ```no_run
68/// # #[cfg(feature = "sqlite")]
69/// # async fn _doc() -> Result<(), sqlx::Error> {
70/// # use sqlx_otel::PoolBuilder;
71/// use sqlx::Executor as _;
72/// use sqlx_otel::QueryAnnotations;
73/// # let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
74///
75/// pool.with_annotations(
76///     QueryAnnotations::new()
77///         .operation("SELECT")
78///         .collection("users"),
79/// )
80/// .fetch_all("SELECT * FROM users")
81/// .await?;
82/// # Ok(())
83/// # }
84/// ```
85#[derive(Debug, Clone, Default, PartialEq, Eq)]
86pub struct QueryAnnotations {
87    /// `db.operation.name` – the database operation (e.g. `"SELECT"`, `"INSERT"`).
88    pub(crate) operation: Option<String>,
89    /// `db.collection.name` – the target table or collection (e.g. `"users"`).
90    pub(crate) collection: Option<String>,
91    /// `db.query.summary` – a low-cardinality summary of the query (e.g. `"SELECT users"`).
92    pub(crate) query_summary: Option<String>,
93    /// `db.stored_procedure.name` – the name of a stored procedure being called.
94    pub(crate) stored_procedure: Option<String>,
95}
96
97impl QueryAnnotations {
98    /// Create a new, empty set of annotations. All fields default to `None`.
99    #[must_use]
100    pub fn new() -> Self {
101        Self::default()
102    }
103
104    /// Set the `db.operation.name` attribute – the database operation being performed
105    /// (e.g. `"SELECT"`, `"INSERT"`, `"findAndModify"`).
106    ///
107    /// The OpenTelemetry semantic conventions require this value to be low cardinality,
108    /// since it is used to construct span names when [`query_summary`](Self::query_summary)
109    /// is not set. Callers who cannot guarantee low cardinality should set
110    /// `query_summary` instead – the library uses that path without a low-cardinality
111    /// assumption.
112    #[must_use]
113    pub fn operation(mut self, operation: impl Into<String>) -> Self {
114        self.operation = Some(operation.into());
115        self
116    }
117
118    /// Set the `db.collection.name` attribute – the table or collection being operated on
119    /// (e.g. `"users"`, `"orders"`).
120    #[must_use]
121    pub fn collection(mut self, collection: impl Into<String>) -> Self {
122        self.collection = Some(collection.into());
123        self
124    }
125
126    /// Set the `db.query.summary` attribute – a low-cardinality summary of the query
127    /// (e.g. `"SELECT users"`, `"INSERT orders"`).
128    ///
129    /// When set, this value also drives the span name (level 1 of the OpenTelemetry
130    /// database span name hierarchy), overriding the `{operation} {collection}`
131    /// synthesis. Cardinality control is the caller's responsibility – a high-cardinality
132    /// summary will produce high-cardinality span names.
133    #[must_use]
134    pub fn query_summary(mut self, summary: impl Into<String>) -> Self {
135        self.query_summary = Some(summary.into());
136        self
137    }
138
139    /// Set the `db.stored_procedure.name` attribute – the name of a stored procedure
140    /// being called (e.g. `"get_user"`, `"sp_update_orders"`).
141    #[must_use]
142    pub fn stored_procedure(mut self, name: impl Into<String>) -> Self {
143        self.stored_procedure = Some(name.into());
144        self
145    }
146}
147
148/// A shared-reference annotation wrapper that carries per-query attributes alongside a
149/// borrowed executor. Returned by [`Pool::with_annotations`](crate::Pool::with_annotations)
150/// and [`Pool::with_operation`](crate::Pool::with_operation).
151///
152/// Implements `sqlx::Executor` with the same instrumentation as the underlying type, but
153/// with annotation values threaded through to span creation.
154pub struct Annotated<'a, E> {
155    pub(crate) inner: &'a E,
156    pub(crate) annotations: QueryAnnotations,
157    pub(crate) state: SharedState,
158}
159
160impl<E: std::fmt::Debug> std::fmt::Debug for Annotated<'_, E> {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        f.debug_struct("Annotated")
163            .field("annotations", &self.annotations)
164            .finish_non_exhaustive()
165    }
166}
167
168/// A mutable-reference annotation wrapper that carries per-query attributes alongside a
169/// mutably borrowed executor. Returned by
170/// [`PoolConnection::with_annotations`](crate::PoolConnection::with_annotations),
171/// [`Transaction::with_annotations`](crate::Transaction::with_annotations), and their
172/// `with_operation` shorthands.
173///
174/// Implements `sqlx::Executor` with the same instrumentation as the underlying type, but
175/// with annotation values threaded through to span creation.
176pub struct AnnotatedMut<'a, E> {
177    pub(crate) inner: &'a mut E,
178    pub(crate) annotations: QueryAnnotations,
179    pub(crate) state: SharedState,
180}
181
182impl<E: std::fmt::Debug> std::fmt::Debug for AnnotatedMut<'_, E> {
183    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
184        f.debug_struct("AnnotatedMut")
185            .field("annotations", &self.annotations)
186            .finish_non_exhaustive()
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193
194    /// Each setter sets exactly its own field and leaves the others untouched. We test
195    /// every permutation (2^4 = 16) including the "none" and "all" cases.
196    #[test]
197    fn setter_permutations() {
198        type Setter = fn(QueryAnnotations) -> QueryAnnotations;
199        type Getter = fn(&QueryAnnotations) -> Option<&str>;
200
201        let fields: &[(&str, Setter, Getter)] = &[
202            (
203                "operation",
204                |a| a.operation("OP"),
205                |a| a.operation.as_deref(),
206            ),
207            (
208                "collection",
209                |a| a.collection("COLL"),
210                |a| a.collection.as_deref(),
211            ),
212            (
213                "query_summary",
214                |a| a.query_summary("SUM"),
215                |a| a.query_summary.as_deref(),
216            ),
217            (
218                "stored_procedure",
219                |a| a.stored_procedure("SP"),
220                |a| a.stored_procedure.as_deref(),
221            ),
222        ];
223
224        for mask in 0u8..16 {
225            let mut ann = QueryAnnotations::new();
226            for (i, &(_, setter, _)) in fields.iter().enumerate() {
227                if mask & (1 << i) != 0 {
228                    ann = setter(ann);
229                }
230            }
231            for (i, &(name, _, getter)) in fields.iter().enumerate() {
232                if mask & (1 << i) != 0 {
233                    assert!(
234                        getter(&ann).is_some(),
235                        "{name} should be Some for mask {mask:#06b}"
236                    );
237                } else {
238                    assert!(
239                        getter(&ann).is_none(),
240                        "{name} should be None for mask {mask:#06b}"
241                    );
242                }
243            }
244        }
245    }
246
247    #[test]
248    fn clone_produces_independent_copy() {
249        let original = QueryAnnotations::new()
250            .operation("SELECT")
251            .collection("users");
252        let cloned = original.clone();
253        let modified = original.query_summary("SELECT users");
254        assert_eq!(cloned.query_summary, None);
255        assert_eq!(modified.query_summary.as_deref(), Some("SELECT users"));
256    }
257
258    #[test]
259    fn debug_impl_is_non_empty() {
260        let ann = QueryAnnotations::new().operation("SELECT");
261        let debug = format!("{ann:?}");
262        assert!(debug.contains("SELECT"));
263    }
264
265    fn test_state() -> SharedState {
266        use std::sync::Arc;
267
268        use crate::attributes::{ConnectionAttributes, QueryTextMode};
269        use crate::metrics::Metrics;
270
271        SharedState {
272            attrs: Arc::new(ConnectionAttributes {
273                system: "sqlite",
274                host: None,
275                port: None,
276                namespace: None,
277                network_peer_address: None,
278                network_peer_port: None,
279                query_text_mode: QueryTextMode::Off,
280            }),
281            metrics: Arc::new(Metrics::new()),
282        }
283    }
284
285    #[test]
286    fn annotated_debug() {
287        let inner = "pool";
288        let wrapper = Annotated {
289            inner: &inner,
290            annotations: QueryAnnotations::new().operation("SELECT"),
291            state: test_state(),
292        };
293        let debug = format!("{wrapper:?}");
294        assert!(debug.contains("Annotated"));
295        assert!(debug.contains("SELECT"));
296    }
297
298    #[test]
299    fn annotated_mut_debug() {
300        let mut inner = "conn";
301        let wrapper = AnnotatedMut {
302            inner: &mut inner,
303            annotations: QueryAnnotations::new().collection("users"),
304            state: test_state(),
305        };
306        let debug = format!("{wrapper:?}");
307        assert!(debug.contains("AnnotatedMut"));
308        assert!(debug.contains("users"));
309    }
310
311    use proptest::prelude::*;
312
313    proptest! {
314        #![proptest_config(ProptestConfig::with_cases(128))]
315
316        /// Calling a setter twice in succession leaves the field equal to the second
317        /// value: each setter is unconditional and overwrites whatever was there.
318        #[test]
319        fn operation_last_write_wins(a in ".{0,64}", b in ".{0,64}") {
320            let ann = QueryAnnotations::new().operation(a).operation(b.clone());
321            prop_assert_eq!(ann.operation.as_deref(), Some(b.as_str()));
322        }
323
324        #[test]
325        fn collection_last_write_wins(a in ".{0,64}", b in ".{0,64}") {
326            let ann = QueryAnnotations::new().collection(a).collection(b.clone());
327            prop_assert_eq!(ann.collection.as_deref(), Some(b.as_str()));
328        }
329
330        #[test]
331        fn query_summary_last_write_wins(a in ".{0,64}", b in ".{0,64}") {
332            let ann = QueryAnnotations::new().query_summary(a).query_summary(b.clone());
333            prop_assert_eq!(ann.query_summary.as_deref(), Some(b.as_str()));
334        }
335
336        #[test]
337        fn stored_procedure_last_write_wins(a in ".{0,64}", b in ".{0,64}") {
338            let ann = QueryAnnotations::new().stored_procedure(a).stored_procedure(b.clone());
339            prop_assert_eq!(ann.stored_procedure.as_deref(), Some(b.as_str()));
340        }
341
342        /// Setting one field never alters any of the others. Complements the
343        /// example-based `setter_permutations` test by stressing arbitrary unicode.
344        #[test]
345        fn operation_does_not_affect_other_fields(s in ".{0,64}") {
346            let ann = QueryAnnotations::new().operation(s);
347            prop_assert!(ann.collection.is_none());
348            prop_assert!(ann.query_summary.is_none());
349            prop_assert!(ann.stored_procedure.is_none());
350        }
351
352        #[test]
353        fn collection_does_not_affect_other_fields(s in ".{0,64}") {
354            let ann = QueryAnnotations::new().collection(s);
355            prop_assert!(ann.operation.is_none());
356            prop_assert!(ann.query_summary.is_none());
357            prop_assert!(ann.stored_procedure.is_none());
358        }
359
360        #[test]
361        fn query_summary_does_not_affect_other_fields(s in ".{0,64}") {
362            let ann = QueryAnnotations::new().query_summary(s);
363            prop_assert!(ann.operation.is_none());
364            prop_assert!(ann.collection.is_none());
365            prop_assert!(ann.stored_procedure.is_none());
366        }
367
368        #[test]
369        fn stored_procedure_does_not_affect_other_fields(s in ".{0,64}") {
370            let ann = QueryAnnotations::new().stored_procedure(s);
371            prop_assert!(ann.operation.is_none());
372            prop_assert!(ann.collection.is_none());
373            prop_assert!(ann.query_summary.is_none());
374        }
375
376        /// All four setters accept arbitrary unicode without panicking, including null
377        /// bytes, surrogate-adjacent code points, and zero-length input.
378        #[test]
379        fn no_panic_setting_all_fields(
380            op in any::<String>(),
381            coll in any::<String>(),
382            summary in any::<String>(),
383            sp in any::<String>(),
384        ) {
385            let _ann = QueryAnnotations::new()
386                .operation(op)
387                .collection(coll)
388                .query_summary(summary)
389                .stored_procedure(sp);
390        }
391
392        /// `Clone` produces a value structurally equal to the original. Combined with
393        /// the `String`-backed fields, this means cloning is a deep copy.
394        #[test]
395        fn clone_equals_original(
396            op in proptest::option::of(".{0,64}"),
397            coll in proptest::option::of(".{0,64}"),
398            summary in proptest::option::of(".{0,64}"),
399            sp in proptest::option::of(".{0,64}"),
400        ) {
401            let mut ann = QueryAnnotations::new();
402            if let Some(s) = op { ann = ann.operation(s); }
403            if let Some(s) = coll { ann = ann.collection(s); }
404            if let Some(s) = summary { ann = ann.query_summary(s); }
405            if let Some(s) = sp { ann = ann.stored_procedure(s); }
406            let cloned = ann.clone();
407            prop_assert_eq!(ann, cloned);
408        }
409    }
410}
sqlx_otel/annotations.rs

sqlx_otel/
annotations.rs
