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 to [`Pool::with_annotations`](crate::Pool::with_annotations),
53/// [`PoolConnection::with_annotations`](crate::PoolConnection::with_annotations), or
54/// [`Transaction::with_annotations`](crate::Transaction::with_annotations).
55///
56/// # Example
57///
58/// ```ignore
59/// pool.with_annotations(
60/// QueryAnnotations::new()
61/// .operation("SELECT")
62/// .collection("users"))
63/// .fetch_all("SELECT * FROM users")
64/// .await?;
65/// ```
66#[derive(Debug, Clone, Default, PartialEq, Eq)]
67pub struct QueryAnnotations {
68 /// `db.operation.name` – the database operation (e.g. `"SELECT"`, `"INSERT"`).
69 pub(crate) operation: Option<String>,
70 /// `db.collection.name` – the target table or collection (e.g. `"users"`).
71 pub(crate) collection: Option<String>,
72 /// `db.query.summary` – a low-cardinality summary of the query (e.g. `"SELECT users"`).
73 pub(crate) query_summary: Option<String>,
74 /// `db.stored_procedure.name` – the name of a stored procedure being called.
75 pub(crate) stored_procedure: Option<String>,
76}
77
78impl QueryAnnotations {
79 /// Create a new, empty set of annotations. All fields default to `None`.
80 #[must_use]
81 pub fn new() -> Self {
82 Self::default()
83 }
84
85 /// Set the `db.operation.name` attribute – the database operation being performed
86 /// (e.g. `"SELECT"`, `"INSERT"`, `"findAndModify"`).
87 ///
88 /// The OpenTelemetry semantic conventions require this value to be low cardinality,
89 /// since it is used to construct span names when [`query_summary`](Self::query_summary)
90 /// is not set. Callers who cannot guarantee low cardinality should set
91 /// `query_summary` instead – the library uses that path without a low-cardinality
92 /// assumption.
93 #[must_use]
94 pub fn operation(mut self, operation: impl Into<String>) -> Self {
95 self.operation = Some(operation.into());
96 self
97 }
98
99 /// Set the `db.collection.name` attribute – the table or collection being operated on
100 /// (e.g. `"users"`, `"orders"`).
101 #[must_use]
102 pub fn collection(mut self, collection: impl Into<String>) -> Self {
103 self.collection = Some(collection.into());
104 self
105 }
106
107 /// Set the `db.query.summary` attribute – a low-cardinality summary of the query
108 /// (e.g. `"SELECT users"`, `"INSERT orders"`).
109 ///
110 /// When set, this value also drives the span name (level 1 of the OpenTelemetry
111 /// database span name hierarchy), overriding the `{operation} {collection}`
112 /// synthesis. Cardinality control is the caller's responsibility – a high-cardinality
113 /// summary will produce high-cardinality span names.
114 #[must_use]
115 pub fn query_summary(mut self, summary: impl Into<String>) -> Self {
116 self.query_summary = Some(summary.into());
117 self
118 }
119
120 /// Set the `db.stored_procedure.name` attribute – the name of a stored procedure
121 /// being called (e.g. `"get_user"`, `"sp_update_orders"`).
122 #[must_use]
123 pub fn stored_procedure(mut self, name: impl Into<String>) -> Self {
124 self.stored_procedure = Some(name.into());
125 self
126 }
127}
128
129/// A shared-reference annotation wrapper that carries per-query attributes alongside a
130/// borrowed executor. Returned by [`Pool::with_annotations`](crate::Pool::with_annotations)
131/// and [`Pool::with_operation`](crate::Pool::with_operation).
132///
133/// Implements `sqlx::Executor` with the same instrumentation as the underlying type, but
134/// with annotation values threaded through to span creation.
135pub struct Annotated<'a, E> {
136 pub(crate) inner: &'a E,
137 pub(crate) annotations: QueryAnnotations,
138 pub(crate) state: SharedState,
139}
140
141impl<E: std::fmt::Debug> std::fmt::Debug for Annotated<'_, E> {
142 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143 f.debug_struct("Annotated")
144 .field("annotations", &self.annotations)
145 .finish_non_exhaustive()
146 }
147}
148
149/// A mutable-reference annotation wrapper that carries per-query attributes alongside a
150/// mutably borrowed executor. Returned by
151/// [`PoolConnection::with_annotations`](crate::PoolConnection::with_annotations),
152/// [`Transaction::with_annotations`](crate::Transaction::with_annotations), and their
153/// `with_operation` shorthands.
154///
155/// Implements `sqlx::Executor` with the same instrumentation as the underlying type, but
156/// with annotation values threaded through to span creation.
157pub struct AnnotatedMut<'a, E> {
158 pub(crate) inner: &'a mut E,
159 pub(crate) annotations: QueryAnnotations,
160 pub(crate) state: SharedState,
161}
162
163impl<E: std::fmt::Debug> std::fmt::Debug for AnnotatedMut<'_, E> {
164 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165 f.debug_struct("AnnotatedMut")
166 .field("annotations", &self.annotations)
167 .finish_non_exhaustive()
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::*;
174
175 /// Each setter sets exactly its own field and leaves the others untouched. We test
176 /// every permutation (2^4 = 16) including the "none" and "all" cases.
177 #[test]
178 fn setter_permutations() {
179 type Setter = fn(QueryAnnotations) -> QueryAnnotations;
180 type Getter = fn(&QueryAnnotations) -> Option<&str>;
181
182 let fields: &[(&str, Setter, Getter)] = &[
183 (
184 "operation",
185 |a| a.operation("OP"),
186 |a| a.operation.as_deref(),
187 ),
188 (
189 "collection",
190 |a| a.collection("COLL"),
191 |a| a.collection.as_deref(),
192 ),
193 (
194 "query_summary",
195 |a| a.query_summary("SUM"),
196 |a| a.query_summary.as_deref(),
197 ),
198 (
199 "stored_procedure",
200 |a| a.stored_procedure("SP"),
201 |a| a.stored_procedure.as_deref(),
202 ),
203 ];
204
205 for mask in 0u8..16 {
206 let mut ann = QueryAnnotations::new();
207 for (i, &(_, setter, _)) in fields.iter().enumerate() {
208 if mask & (1 << i) != 0 {
209 ann = setter(ann);
210 }
211 }
212 for (i, &(name, _, getter)) in fields.iter().enumerate() {
213 if mask & (1 << i) != 0 {
214 assert!(
215 getter(&ann).is_some(),
216 "{name} should be Some for mask {mask:#06b}"
217 );
218 } else {
219 assert!(
220 getter(&ann).is_none(),
221 "{name} should be None for mask {mask:#06b}"
222 );
223 }
224 }
225 }
226 }
227
228 #[test]
229 fn clone_produces_independent_copy() {
230 let original = QueryAnnotations::new()
231 .operation("SELECT")
232 .collection("users");
233 let cloned = original.clone();
234 let modified = original.query_summary("SELECT users");
235 assert_eq!(cloned.query_summary, None);
236 assert_eq!(modified.query_summary.as_deref(), Some("SELECT users"));
237 }
238
239 #[test]
240 fn debug_impl_is_non_empty() {
241 let ann = QueryAnnotations::new().operation("SELECT");
242 let debug = format!("{ann:?}");
243 assert!(debug.contains("SELECT"));
244 }
245
246 fn test_state() -> SharedState {
247 use std::sync::Arc;
248
249 use crate::attributes::{ConnectionAttributes, QueryTextMode};
250 use crate::metrics::Metrics;
251
252 SharedState {
253 attrs: Arc::new(ConnectionAttributes {
254 system: "sqlite",
255 host: None,
256 port: None,
257 namespace: None,
258 network_peer_address: None,
259 network_peer_port: None,
260 query_text_mode: QueryTextMode::Off,
261 }),
262 metrics: Arc::new(Metrics::new()),
263 }
264 }
265
266 #[test]
267 fn annotated_debug() {
268 let inner = "pool";
269 let wrapper = Annotated {
270 inner: &inner,
271 annotations: QueryAnnotations::new().operation("SELECT"),
272 state: test_state(),
273 };
274 let debug = format!("{wrapper:?}");
275 assert!(debug.contains("Annotated"));
276 assert!(debug.contains("SELECT"));
277 }
278
279 #[test]
280 fn annotated_mut_debug() {
281 let mut inner = "conn";
282 let wrapper = AnnotatedMut {
283 inner: &mut inner,
284 annotations: QueryAnnotations::new().collection("users"),
285 state: test_state(),
286 };
287 let debug = format!("{wrapper:?}");
288 assert!(debug.contains("AnnotatedMut"));
289 assert!(debug.contains("users"));
290 }
291}