sqlx_otel/
query_ext.rs

1//! Query-side annotation surface – mirror of the executor-side
2//! [`with_annotations`](crate::Pool::with_annotations) /
3//! [`with_operation`](crate::Pool::with_operation) methods, but attached to the query
4//! builder produced by [`sqlx::query()`], [`sqlx::query_as()`], and
5//! [`sqlx::query_scalar()`].
6//!
7//! # Choosing between executor-side and query-side
8//!
9//! The two surfaces emit identical telemetry; pick whichever keeps the annotation closer to
10//! the thing it describes. As a rule of thumb:
11//!
12//! - **Query-side** is the default. Per-query attributes describe the *query*, and
13//!   colocating them with the query text keeps both readable.
14//! - **Executor-side** is for cases where one annotation set is reused across many queries
15//!   on the same executor (e.g. a request-scoped logical operation), or where the surface
16//!   is `prepare` / `describe` rather than `execute` / `fetch*`.
17//!
18//! ```no_run
19//! # #[cfg(feature = "sqlite")]
20//! # async fn _doc() -> Result<(), sqlx::Error> {
21//! # use sqlx_otel::PoolBuilder;
22//! use sqlx_otel::{QueryAnnotateExt, QueryAnnotations};
23//! # let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
24//!
25//! sqlx::query("SELECT * FROM users WHERE id = ?")
26//!     .bind(42_i64)
27//!     .with_annotations(
28//!         QueryAnnotations::new()
29//!             .operation("SELECT")
30//!             .collection("users"),
31//!     )
32//!     .execute(&pool)
33//!     .await?;
34//! # Ok(())
35//! # }
36//! ```
37//!
38//! The wrapper [`AnnotatedQuery`] exposes the same `execute`/`fetch*` surface as the inner
39//! query and threads the annotations into span creation by wrapping the executor with the
40//! existing [`Annotated`](Annotated) / [`AnnotatedMut`](AnnotatedMut) executor wrappers.
41//!
42//! # Map and macro queries
43//!
44//! [`Query::map`](Query::map) / [`Query::try_map`](Query::try_map) return
45//! [`sqlx::query::Map<'q, DB, F, A>`](sqlx::query::Map), which is also covered by the
46//! trait. `with_annotations` and `with_operation` may be applied at any of the three
47//! positions on a hand-written `Query::map()` chain – before `bind`, between `bind` and
48//! `map`, or after `map`:
49//!
50//! ```no_run
51//! # #[cfg(feature = "sqlite")]
52//! # async fn _doc() -> Result<(), sqlx::Error> {
53//! # use sqlx_otel::PoolBuilder;
54//! use sqlx::Row as _;
55//! use sqlx_otel::QueryAnnotateExt;
56//! # let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
57//!
58//! sqlx::query("SELECT id FROM users WHERE name = ?")
59//!     .bind("alice")
60//!     .map(|row: sqlx::sqlite::SqliteRow| row.get::<i64, _>("id"))
61//!     .with_operation("SELECT", "users")
62//!     .fetch_one(&pool)
63//!     .await?;
64//! # Ok(())
65//! # }
66//! ```
67//!
68//! The compile-time validated macro forms (`sqlx::query!()`, `sqlx::query_as!()`,
69//! `sqlx::query_scalar!()`) expand to either `Query<'q, DB, _>` (for no-result-column
70//! shapes) or `Map<'q, DB, _, _>` (for any shape that decodes columns). Both are covered:
71//!
72//! ```ignore
73//! sqlx::query_as!(User, "SELECT id, name FROM users WHERE id = ?", 42_i64)
74//!     .with_operation("SELECT", "users")
75//!     .fetch_one(&pool)
76//!     .await?;
77//! ```
78//!
79//! Macro queries can only carry annotations *after* the macro returns – the macro itself
80//! pre-applies `bind` and `try_map`, so positions 1 and 2 are not reachable by the user.
81//! The macro example above stays `ignore` because it requires either an offline `.sqlx/`
82//! cache or a live `DATABASE_URL`.
83
84use futures::stream::BoxStream;
85use sqlx::query::{Map, Query, QueryAs, QueryScalar};
86
87use crate::annotations::{Annotated, AnnotatedMut, QueryAnnotations};
88use crate::database::Database;
89
90/// Sealing module for [`QueryAnnotateExt`]. Downstream crates cannot implement the trait
91/// because the wrapper types it produces (`Annotated` / `AnnotatedMut`) have private fields
92/// that only this crate can construct.
93mod sealed {
94    /// Marker trait that prevents external impls of [`super::QueryAnnotateExt`].
95    pub trait Sealed {}
96}
97
98/// Extension trait that attaches OpenTelemetry per-query annotations to the function-form
99/// `SQLx` query builders ([`sqlx::query()`], [`sqlx::query_as()`],
100/// [`sqlx::query_scalar()`]) and to the [`Map`](sqlx::query::Map) returned by
101/// `Query::map` / `Query::try_map`.
102///
103/// The trait is sealed and cannot be implemented downstream.
104///
105/// # Example
106///
107/// ```no_run
108/// # #[cfg(feature = "sqlite")]
109/// # async fn _doc() -> Result<(), sqlx::Error> {
110/// # use sqlx_otel::PoolBuilder;
111/// use sqlx_otel::QueryAnnotateExt;
112/// # let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
113///
114/// sqlx::query("INSERT INTO orders (user_id) VALUES (?)")
115///     .bind(42_i64)
116///     .with_operation("INSERT", "orders")
117///     .execute(&pool)
118///     .await?;
119/// # Ok(())
120/// # }
121/// ```
122pub trait QueryAnnotateExt: sealed::Sealed + Sized {
123    /// Wrap the query with the given per-query annotations.
124    ///
125    /// Returns an [`AnnotatedQuery`] that exposes the same `execute`/`fetch*` surface as the
126    /// inner query but threads the annotations through to OpenTelemetry span creation.
127    fn with_annotations(self, annotations: QueryAnnotations) -> AnnotatedQuery<Self>;
128
129    /// Shorthand that sets `db.operation.name` and `db.collection.name`.
130    ///
131    /// Equivalent to
132    /// `self.with_annotations(QueryAnnotations::new().operation(op).collection(coll))`.
133    ///
134    /// The returned future is `Send` and may be used inside a `Send`-required async context
135    /// (e.g. `tokio::spawn`, axum handlers, `tower::Service`-bounded futures). The example
136    /// below intentionally swallows the inner `sqlx::Error` via `.ok().flatten()` – the
137    /// point is to exercise the `Send` contract on the spawned future at compile time, not
138    /// to demonstrate error handling. A non-`Send` future would fail to compile here, since
139    /// `tokio::spawn` requires `Send`.
140    ///
141    /// ```no_run
142    /// # #[cfg(feature = "sqlite")]
143    /// # async fn _doc() -> Result<(), Box<dyn std::error::Error>> {
144    /// # use sqlx_otel::PoolBuilder;
145    /// use sqlx_otel::QueryAnnotateExt as _;
146    /// # let pool: sqlx_otel::Pool<sqlx::Sqlite> =
147    /// #     PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
148    ///
149    /// tokio::spawn(async move {
150    ///     let _: Option<(i32,)> = sqlx::query_as("SELECT 1")
151    ///         .with_operation("SELECT", "users")
152    ///         .fetch_optional(&pool)
153    ///         .await
154    ///         .ok()
155    ///         .flatten();
156    /// })
157    /// .await?;
158    /// # Ok(()) }
159    /// ```
160    fn with_operation(
161        self,
162        operation: impl Into<String>,
163        collection: impl Into<String>,
164    ) -> AnnotatedQuery<Self> {
165        self.with_annotations(
166            QueryAnnotations::new()
167                .operation(operation)
168                .collection(collection),
169        )
170    }
171}
172
173impl<DB: sqlx::Database, A> sealed::Sealed for Query<'_, DB, A> {}
174impl<DB: sqlx::Database, A> QueryAnnotateExt for Query<'_, DB, A> {
175    fn with_annotations(self, annotations: QueryAnnotations) -> AnnotatedQuery<Self> {
176        AnnotatedQuery {
177            inner: self,
178            annotations,
179        }
180    }
181}
182
183impl<DB: sqlx::Database, O, A> sealed::Sealed for QueryAs<'_, DB, O, A> {}
184impl<DB: sqlx::Database, O, A> QueryAnnotateExt for QueryAs<'_, DB, O, A> {
185    fn with_annotations(self, annotations: QueryAnnotations) -> AnnotatedQuery<Self> {
186        AnnotatedQuery {
187            inner: self,
188            annotations,
189        }
190    }
191}
192
193impl<DB: sqlx::Database, O, A> sealed::Sealed for QueryScalar<'_, DB, O, A> {}
194impl<DB: sqlx::Database, O, A> QueryAnnotateExt for QueryScalar<'_, DB, O, A> {
195    fn with_annotations(self, annotations: QueryAnnotations) -> AnnotatedQuery<Self> {
196        AnnotatedQuery {
197            inner: self,
198            annotations,
199        }
200    }
201}
202
203impl<DB: sqlx::Database, F, A> sealed::Sealed for Map<'_, DB, F, A> {}
204impl<DB: sqlx::Database, F, A> QueryAnnotateExt for Map<'_, DB, F, A> {
205    fn with_annotations(self, annotations: QueryAnnotations) -> AnnotatedQuery<Self> {
206        AnnotatedQuery {
207            inner: self,
208            annotations,
209        }
210    }
211}
212
213/// A `SQLx` query builder paired with OpenTelemetry per-query annotations.
214///
215/// Produced by [`QueryAnnotateExt::with_annotations`] / [`QueryAnnotateExt::with_operation`].
216/// Each invocation of the executor-driven methods (`execute`, `fetch_all`, etc.) wraps the
217/// executor with the annotations so the resulting span carries them.
218///
219/// The wrapper exposes [`bind`](AnnotatedQuery::bind) so the caller can chain
220/// `.bind(...).with_annotations(...)` or `.with_annotations(...).bind(...)` interchangeably.
221#[must_use = "annotated queries do nothing until you call execute / fetch* on them"]
222pub struct AnnotatedQuery<Q> {
223    inner: Q,
224    annotations: QueryAnnotations,
225}
226
227impl<Q> std::fmt::Debug for AnnotatedQuery<Q> {
228    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
229        f.debug_struct("AnnotatedQuery")
230            .field("annotations", &self.annotations)
231            .finish_non_exhaustive()
232    }
233}
234
235// ---------------------------------------------------------------------------
236// Internal trait: convert an executor argument into the existing Annotated /
237// AnnotatedMut wrappers. Implemented for the same three borrow shapes that the
238// `impl_executor!` macro covers.
239// ---------------------------------------------------------------------------
240
241/// Internal conversion trait. Implemented for the three executor borrow shapes that
242/// `impl_executor!` already supports (`&Pool`, `&mut PoolConnection`, `&mut Transaction`),
243/// each producing the matching [`Annotated`] / [`AnnotatedMut`] wrapper.
244///
245/// The HRTB `for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>` lives on each
246/// individual impl rather than on the trait, to avoid trait-resolution recursion when the
247/// user's executor type is itself constructed via the same HRTB. This bound used to leak into
248/// the `Send` auto-trait inference of `AnnotatedQuery::fetch_*` / `execute` when those methods
249/// were `async fn` – the resulting opaque coroutine carried a region-quantified obligation
250/// that rustc could not always discharge, which broke query-side annotation inside contexts
251/// requiring `Send` (`tokio::spawn`, axum handlers). The fix is in the *callers* of this
252/// trait, not the trait itself: the forwarders on `AnnotatedQuery` are written as
253/// `fn(...) -> impl Future + Send + 'e` returning `SQLx`'s own future directly, so no
254/// coroutine forms on this crate's side and the `Send` check delegates to the underlying
255/// `SQLx` future, which is already `Send`-clean for concrete backends.
256///
257/// Users do not call this trait directly – they pass `&pool`, `&mut conn`, or `&mut tx` to
258/// [`AnnotatedQuery::execute`] / `fetch*` and the trait dispatches internally. The trait is
259/// effectively sealed because only this crate can construct the [`Annotated`] /
260/// [`AnnotatedMut`] wrapper types (their fields are crate-private).
261#[doc(hidden)]
262pub trait IntoAnnotatedExecutor<'e, DB: Database> {
263    /// The annotated executor wrapper type produced by [`into_annotated`](Self::into_annotated).
264    type Wrapper: sqlx::Executor<'e, Database = DB>;
265
266    /// Consume `self` together with the per-query annotations to produce the wrapper.
267    fn into_annotated(self, annotations: QueryAnnotations) -> Self::Wrapper;
268}
269
270impl<'e, DB> IntoAnnotatedExecutor<'e, DB> for &'e crate::Pool<DB>
271where
272    DB: Database,
273    for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
274{
275    type Wrapper = Annotated<'e, crate::Pool<DB>>;
276
277    fn into_annotated(self, annotations: QueryAnnotations) -> Self::Wrapper {
278        Annotated {
279            inner: self,
280            annotations,
281            state: self.state.clone(),
282        }
283    }
284}
285
286impl<'e, DB> IntoAnnotatedExecutor<'e, DB> for &'e mut crate::PoolConnection<DB>
287where
288    DB: Database,
289    for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
290{
291    type Wrapper = AnnotatedMut<'e, crate::PoolConnection<DB>>;
292
293    fn into_annotated(self, annotations: QueryAnnotations) -> Self::Wrapper {
294        AnnotatedMut {
295            state: self.state.clone(),
296            annotations,
297            inner: self,
298        }
299    }
300}
301
302impl<'e, 'tx, DB> IntoAnnotatedExecutor<'e, DB> for &'e mut crate::Transaction<'tx, DB>
303where
304    DB: Database,
305    for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
306{
307    type Wrapper = AnnotatedMut<'e, crate::Transaction<'tx, DB>>;
308
309    fn into_annotated(self, annotations: QueryAnnotations) -> Self::Wrapper {
310        AnnotatedMut {
311            state: self.state.clone(),
312            annotations,
313            inner: self,
314        }
315    }
316}
317
318// ---------------------------------------------------------------------------
319// Forwarder macros: emit method bodies inside a caller-supplied `impl` block.
320//
321// The macros expand to method items only (no `impl` header), so each call site
322// declares its own generics and where-clauses. This avoids the local-ambiguity
323// that arises when the macro tries to consume an `impl<...>` header via `tt`
324// repetition, while still removing the per-builder duplication.
325//
326// Symbols referenced in the bodies (`Self`, `DB`, `A`, `O`, `QueryAnnotations`,
327// `IntoAnnotatedExecutor`, `BoxStream`) resolve in the expansion context.
328// ---------------------------------------------------------------------------
329
330/// Emit the methods that every `AnnotatedQuery<Inner>` impl block has in common:
331/// `with_annotations`, `with_operation`, and the five `fetch*` forwarders.
332///
333/// Parameters:
334/// * `row` – the row type of the inner builder (`DB::Row` for `Query`, `O` for
335///   `QueryAs` and `QueryScalar`).
336/// * `extra_bounds` – additional lifetime bounds threaded onto each fetch-forwarder's
337///   where-clause (`DB: 'e, O: 'e,` for `QueryAs` / `QueryScalar`; empty for `Query`).
338macro_rules! impl_annotated_query_fetch_forwarders {
339    (row = $row:ty, extra_bounds = ($($extra_bounds:tt)*)) => {
340        /// Replace the wrapper's annotations. Last-call-wins – any previously set annotations
341        /// on this wrapper are discarded.
342        pub fn with_annotations(mut self, annotations: QueryAnnotations) -> Self {
343            self.annotations = annotations;
344            self
345        }
346
347        /// Shorthand replacement that sets `db.operation.name` and `db.collection.name`.
348        /// Discards any previously set annotations on this wrapper.
349        pub fn with_operation(
350            self,
351            operation: impl Into<String>,
352            collection: impl Into<String>,
353        ) -> Self {
354            self.with_annotations(
355                QueryAnnotations::new()
356                    .operation(operation)
357                    .collection(collection),
358            )
359        }
360
361        /// Stream the resulting rows.
362        pub fn fetch<'e, E>(self, executor: E) -> BoxStream<'e, Result<$row, sqlx::Error>>
363        where
364            'q: 'e,
365            A: 'e,
366            $($extra_bounds)*
367            E: 'e + IntoAnnotatedExecutor<'e, DB>,
368        {
369            let wrapper = executor.into_annotated(self.annotations);
370            self.inner.fetch(wrapper)
371        }
372
373        /// Stream a mix of `QueryResult`s and rows for multi-statement queries.
374        ///
375        /// `fetch_many` is `#[deprecated]` in `SQLx` 0.8 but kept for parity with the
376        /// executor-side surface.
377        #[allow(deprecated, clippy::type_complexity)]
378        pub fn fetch_many<'e, E>(
379            self,
380            executor: E,
381        ) -> BoxStream<'e, Result<sqlx::Either<DB::QueryResult, $row>, sqlx::Error>>
382        where
383            'q: 'e,
384            A: 'e,
385            $($extra_bounds)*
386            E: 'e + IntoAnnotatedExecutor<'e, DB>,
387        {
388            let wrapper = executor.into_annotated(self.annotations);
389            self.inner.fetch_many(wrapper)
390        }
391
392        /// Collect every row into a `Vec`.
393        ///
394        /// # Errors
395        ///
396        /// Returns any [`sqlx::Error`] surfaced by the underlying driver, including row
397        /// decoding errors.
398        pub fn fetch_all<'e, E>(
399            self,
400            executor: E,
401        ) -> impl 'e + Send + std::future::Future<Output = Result<Vec<$row>, sqlx::Error>>
402        where
403            'q: 'e,
404            A: 'e,
405            $($extra_bounds)*
406            E: 'e + IntoAnnotatedExecutor<'e, DB>,
407        {
408            let wrapper = executor.into_annotated(self.annotations);
409            self.inner.fetch_all(wrapper)
410        }
411
412        /// Return exactly one row, erroring if none or more than one.
413        ///
414        /// # Errors
415        ///
416        /// Returns [`sqlx::Error::RowNotFound`] when the result set is empty, or any other
417        /// [`sqlx::Error`] surfaced by the underlying driver.
418        pub fn fetch_one<'e, E>(
419            self,
420            executor: E,
421        ) -> impl 'e + Send + std::future::Future<Output = Result<$row, sqlx::Error>>
422        where
423            'q: 'e,
424            A: 'e,
425            $($extra_bounds)*
426            E: 'e + IntoAnnotatedExecutor<'e, DB>,
427        {
428            let wrapper = executor.into_annotated(self.annotations);
429            self.inner.fetch_one(wrapper)
430        }
431
432        /// Return at most one row.
433        ///
434        /// # Errors
435        ///
436        /// Returns any [`sqlx::Error`] surfaced by the underlying driver.
437        pub fn fetch_optional<'e, E>(
438            self,
439            executor: E,
440        ) -> impl 'e + Send + std::future::Future<Output = Result<Option<$row>, sqlx::Error>>
441        where
442            'q: 'e,
443            A: 'e,
444            $($extra_bounds)*
445            E: 'e + IntoAnnotatedExecutor<'e, DB>,
446        {
447            let wrapper = executor.into_annotated(self.annotations);
448            self.inner.fetch_optional(wrapper)
449        }
450    };
451}
452
453/// Emit the `bind` method on a specialised impl block.
454///
455/// `SQLx` restricts `bind` to the default `<DB as Database>::Arguments<'q>` parameter, so
456/// each builder family has its own dedicated impl block keyed on the default arguments.
457macro_rules! impl_annotated_query_bind {
458    () => {
459        /// Append a parameter binding, forwarding to the inner query. Mirrors the inner
460        /// builder's own `bind` method.
461        pub fn bind<T>(mut self, value: T) -> Self
462        where
463            T: 'q + sqlx::Encode<'q, DB> + sqlx::Type<DB>,
464        {
465            self.inner = self.inner.bind(value);
466            self
467        }
468    };
469}
470
471// --- AnnotatedQuery<Query<'q, DB, A>> --------------------------------------
472
473impl<'q, DB, A> AnnotatedQuery<Query<'q, DB, A>>
474where
475    DB: Database,
476    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
477{
478    impl_annotated_query_fetch_forwarders!(row = DB::Row, extra_bounds = ());
479
480    /// Execute the query and return the number of rows affected. Wraps the executor with
481    /// the carried annotations so the resulting span is annotated.
482    ///
483    /// # Errors
484    ///
485    /// Returns any [`sqlx::Error`] surfaced by the underlying driver.
486    pub fn execute<'e, E>(
487        self,
488        executor: E,
489    ) -> impl 'e + Send + std::future::Future<Output = Result<DB::QueryResult, sqlx::Error>>
490    where
491        'q: 'e,
492        A: 'e,
493        E: 'e + IntoAnnotatedExecutor<'e, DB>,
494    {
495        let wrapper = executor.into_annotated(self.annotations);
496        self.inner.execute(wrapper)
497    }
498
499    /// Execute multiple statements separated by `;` and return their results as a stream.
500    ///
501    /// `execute_many` is `#[deprecated]` in `SQLx` 0.8 but kept here for parity with the
502    /// existing executor-side surface. Only `Query` exposes this method – `QueryAs`,
503    /// `QueryScalar`, and `Map` have no `execute_many` upstream.
504    #[allow(deprecated)]
505    pub fn execute_many<'e, E>(
506        self,
507        executor: E,
508    ) -> impl 'e + Send + std::future::Future<Output = BoxStream<'e, Result<DB::QueryResult, sqlx::Error>>>
509    where
510        'q: 'e,
511        A: 'e,
512        E: 'e + IntoAnnotatedExecutor<'e, DB>,
513    {
514        let wrapper = executor.into_annotated(self.annotations);
515        self.inner.execute_many(wrapper)
516    }
517
518    /// Map each row to another type. Mirrors [`sqlx::query::Query::map`] and carries the
519    /// existing annotations forward unchanged onto the resulting `AnnotatedQuery<Map<...>>`,
520    /// so `with_annotations` can be applied either before or after `.map()`.
521    #[allow(clippy::type_complexity)]
522    pub fn map<F, O>(
523        self,
524        f: F,
525    ) -> AnnotatedQuery<Map<'q, DB, impl FnMut(DB::Row) -> Result<O, sqlx::Error> + Send, A>>
526    where
527        F: FnMut(DB::Row) -> O + Send,
528        O: Unpin,
529    {
530        AnnotatedQuery {
531            inner: self.inner.map(f),
532            annotations: self.annotations,
533        }
534    }
535
536    /// Map each row to a `Result`. Mirrors [`sqlx::query::Query::try_map`] and carries the
537    /// existing annotations forward unchanged.
538    pub fn try_map<F, O>(self, f: F) -> AnnotatedQuery<Map<'q, DB, F, A>>
539    where
540        F: FnMut(DB::Row) -> Result<O, sqlx::Error> + Send,
541        O: Unpin,
542    {
543        AnnotatedQuery {
544            inner: self.inner.try_map(f),
545            annotations: self.annotations,
546        }
547    }
548}
549
550impl<'q, DB> AnnotatedQuery<Query<'q, DB, <DB as sqlx::Database>::Arguments<'q>>>
551where
552    DB: sqlx::Database,
553{
554    impl_annotated_query_bind!();
555}
556
557// --- AnnotatedQuery<QueryAs<'q, DB, O, A>> ---------------------------------
558
559impl<'q, DB, O, A> AnnotatedQuery<QueryAs<'q, DB, O, A>>
560where
561    DB: Database,
562    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
563    O: Send + Unpin + for<'r> sqlx::FromRow<'r, DB::Row>,
564{
565    impl_annotated_query_fetch_forwarders!(row = O, extra_bounds = (DB: 'e, O: 'e,));
566}
567
568impl<'q, DB, O> AnnotatedQuery<QueryAs<'q, DB, O, <DB as sqlx::Database>::Arguments<'q>>>
569where
570    DB: sqlx::Database,
571{
572    impl_annotated_query_bind!();
573}
574
575// --- AnnotatedQuery<QueryScalar<'q, DB, O, A>> -----------------------------
576
577impl<'q, DB, O, A> AnnotatedQuery<QueryScalar<'q, DB, O, A>>
578where
579    DB: Database,
580    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
581    O: Send + Unpin,
582    (O,): Send + Unpin + for<'r> sqlx::FromRow<'r, DB::Row>,
583{
584    impl_annotated_query_fetch_forwarders!(row = O, extra_bounds = (DB: 'e, O: 'e,));
585}
586
587impl<'q, DB, O> AnnotatedQuery<QueryScalar<'q, DB, O, <DB as sqlx::Database>::Arguments<'q>>>
588where
589    DB: sqlx::Database,
590{
591    impl_annotated_query_bind!();
592}
593
594// --- AnnotatedQuery<Map<'q, DB, F, A>> -------------------------------------
595
596impl<'q, DB, F, A, O> AnnotatedQuery<Map<'q, DB, F, A>>
597where
598    DB: Database,
599    F: FnMut(DB::Row) -> Result<O, sqlx::Error> + Send,
600    O: Send + Unpin,
601    A: 'q + Send + sqlx::IntoArguments<'q, DB>,
602{
603    impl_annotated_query_fetch_forwarders!(row = O, extra_bounds = (DB: 'e, F: 'e, O: 'e,));
604
605    /// Compose a further mapping on top of this annotated map. Mirrors
606    /// [`sqlx::query::Map::map`] (which itself composes via `f(row).and_then(&mut g)`) and
607    /// preserves the existing annotations on the wrapper.
608    #[allow(clippy::type_complexity)]
609    pub fn map<G, P>(
610        self,
611        g: G,
612    ) -> AnnotatedQuery<Map<'q, DB, impl FnMut(DB::Row) -> Result<P, sqlx::Error> + Send, A>>
613    where
614        G: FnMut(O) -> P + Send,
615        P: Unpin,
616    {
617        AnnotatedQuery {
618            inner: self.inner.map(g),
619            annotations: self.annotations,
620        }
621    }
622
623    /// Fallible variant of [`map`](Self::map). Mirrors [`sqlx::query::Map::try_map`] and
624    /// preserves the existing annotations on the wrapper.
625    #[allow(clippy::type_complexity)]
626    pub fn try_map<G, P>(
627        self,
628        g: G,
629    ) -> AnnotatedQuery<Map<'q, DB, impl FnMut(DB::Row) -> Result<P, sqlx::Error> + Send, A>>
630    where
631        G: FnMut(O) -> Result<P, sqlx::Error> + Send,
632        P: Unpin,
633    {
634        AnnotatedQuery {
635            inner: self.inner.try_map(g),
636            annotations: self.annotations,
637        }
638    }
639}
640
641#[cfg(all(test, feature = "sqlite"))]
642mod tests {
643    use sqlx::Execute as _;
644    use sqlx::Sqlite;
645
646    use super::*;
647
648    // --- query() / query_as() / query_scalar() --------------------
649
650    #[test]
651    fn with_annotations_replaces_previous() {
652        let q = sqlx::query::<Sqlite>("SELECT 1")
653            .with_annotations(QueryAnnotations::new().operation("FIRST"))
654            .with_annotations(QueryAnnotations::new().operation("SECOND"));
655        assert_eq!(q.annotations.operation.as_deref(), Some("SECOND"));
656    }
657
658    #[test]
659    fn with_operation_sets_both_fields() {
660        let q = sqlx::query::<Sqlite>("SELECT 1").with_operation("SELECT", "users");
661        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
662        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
663        assert!(q.annotations.query_summary.is_none());
664        assert!(q.annotations.stored_procedure.is_none());
665    }
666
667    #[test]
668    fn with_operation_replaces_previous_annotations() {
669        // Documents the last-call-wins behaviour: with_operation discards earlier fields.
670        let q = sqlx::query::<Sqlite>("SELECT 1")
671            .with_annotations(QueryAnnotations::new().query_summary("legacy summary"))
672            .with_operation("SELECT", "users");
673        assert!(
674            q.annotations.query_summary.is_none(),
675            "with_operation must replace, not merge"
676        );
677    }
678
679    #[test]
680    fn debug_impl_includes_annotations() {
681        let q = sqlx::query::<Sqlite>("SELECT 1")
682            .with_annotations(QueryAnnotations::new().operation("DEBUG_OP"));
683        let debug = format!("{q:?}");
684        assert!(debug.contains("AnnotatedQuery"));
685        assert!(debug.contains("DEBUG_OP"));
686    }
687
688    #[test]
689    fn bind_then_with_annotations_compose() {
690        let q = sqlx::query::<Sqlite>("SELECT ?1, ?2")
691            .bind(1_i32)
692            .with_annotations(QueryAnnotations::new().operation("SELECT"));
693        assert_eq!(q.inner.sql(), "SELECT ?1, ?2");
694        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
695    }
696
697    #[test]
698    fn with_annotations_first_then_bind_compose() {
699        let q = sqlx::query::<Sqlite>("SELECT ?1, ?2")
700            .with_annotations(QueryAnnotations::new().operation("SELECT"))
701            .bind(1_i32);
702        assert_eq!(q.inner.sql(), "SELECT ?1, ?2");
703        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
704    }
705
706    #[test]
707    fn query_as_supports_with_annotations() {
708        let q = sqlx::query_as::<Sqlite, (i32,)>("SELECT 1").with_operation("SELECT", "users");
709        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
710        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
711    }
712
713    #[test]
714    fn query_as_wrapper_with_annotations_replaces() {
715        // Exercise the inherent `with_annotations` on AnnotatedQuery<QueryAs<_>>.
716        let q = sqlx::query_as::<Sqlite, (i32,)>("SELECT 1")
717            .with_annotations(QueryAnnotations::new().operation("FIRST"))
718            .with_annotations(QueryAnnotations::new().operation("SECOND"));
719        assert_eq!(q.annotations.operation.as_deref(), Some("SECOND"));
720    }
721
722    #[test]
723    fn query_as_wrapper_with_operation_chains_on_wrapper() {
724        // First call uses the trait method (on QueryAs); second uses the inherent method
725        // on AnnotatedQuery, which exercises the wrapper-level shorthand.
726        let q = sqlx::query_as::<Sqlite, (i32,)>("SELECT 1")
727            .with_annotations(QueryAnnotations::new().query_summary("legacy"))
728            .with_operation("SELECT", "users");
729        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
730        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
731        assert!(q.annotations.query_summary.is_none());
732    }
733
734    #[test]
735    fn query_as_wrapper_bind_after_annotations() {
736        let q = sqlx::query_as::<Sqlite, (i32,)>("SELECT ?1")
737            .with_annotations(QueryAnnotations::new().operation("SELECT"))
738            .bind(7_i32);
739        assert_eq!(q.inner.sql(), "SELECT ?1");
740        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
741    }
742
743    #[test]
744    fn query_scalar_supports_with_annotations() {
745        let q = sqlx::query_scalar::<Sqlite, i32>("SELECT 1").with_operation("SELECT", "users");
746        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
747        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
748    }
749
750    #[test]
751    fn query_scalar_wrapper_with_annotations_replaces() {
752        let q = sqlx::query_scalar::<Sqlite, i32>("SELECT 1")
753            .with_annotations(QueryAnnotations::new().operation("FIRST"))
754            .with_annotations(QueryAnnotations::new().operation("SECOND"));
755        assert_eq!(q.annotations.operation.as_deref(), Some("SECOND"));
756    }
757
758    #[test]
759    fn query_scalar_wrapper_with_operation_chains_on_wrapper() {
760        let q = sqlx::query_scalar::<Sqlite, i32>("SELECT 1")
761            .with_annotations(QueryAnnotations::new().query_summary("legacy"))
762            .with_operation("SELECT", "users");
763        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
764        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
765        assert!(q.annotations.query_summary.is_none());
766    }
767
768    #[test]
769    fn query_scalar_wrapper_bind_after_annotations() {
770        let q = sqlx::query_scalar::<Sqlite, i32>("SELECT ?1")
771            .with_annotations(QueryAnnotations::new().operation("SELECT"))
772            .bind(7_i32);
773        assert_eq!(q.inner.sql(), "SELECT ?1");
774        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
775    }
776
777    // --- Map / Query::map / Query::try_map composition --------------------
778
779    #[test]
780    fn query_with_annotations_map_preserves_annotations() {
781        // Position 1: annotate before `.map()`. Annotations must survive the wrap.
782        let q = sqlx::query::<Sqlite>("SELECT 1")
783            .with_annotations(QueryAnnotations::new().operation("SELECT"))
784            .map(|_row: sqlx::sqlite::SqliteRow| 42_i64);
785        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
786    }
787
788    #[test]
789    fn query_bind_with_annotations_map_preserves_annotations() {
790        // Position 2: annotate between `.bind()` and `.map()`.
791        let q = sqlx::query::<Sqlite>("SELECT ?1")
792            .bind(1_i32)
793            .with_annotations(QueryAnnotations::new().operation("SELECT"))
794            .map(|_row: sqlx::sqlite::SqliteRow| 42_i64);
795        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
796    }
797
798    #[test]
799    fn query_map_with_annotations_replaces_previous() {
800        // Position 3 + last-call-wins on the new `Map` wrapper.
801        let q = sqlx::query::<Sqlite>("SELECT 1")
802            .map(|_row: sqlx::sqlite::SqliteRow| 42_i64)
803            .with_annotations(QueryAnnotations::new().operation("FIRST"))
804            .with_annotations(QueryAnnotations::new().operation("SECOND"));
805        assert_eq!(q.annotations.operation.as_deref(), Some("SECOND"));
806    }
807
808    #[test]
809    fn query_try_map_with_annotations_compose() {
810        // `try_map` stores `F` directly; sanity-check the non-opaque closure branch.
811        let q = sqlx::query::<Sqlite>("SELECT 1")
812            .try_map(|_row: sqlx::sqlite::SqliteRow| Ok::<_, sqlx::Error>(42_i64))
813            .with_operation("SELECT", "users");
814        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
815        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
816    }
817
818    #[test]
819    fn annotated_query_try_map_preserves_annotations() {
820        // Exercise `AnnotatedQuery<Query>::try_map` (the wrapper's own method, not sqlx's).
821        // Position-1-then-fallible-mapper.
822        let q = sqlx::query::<Sqlite>("SELECT 1")
823            .with_annotations(QueryAnnotations::new().operation("SELECT"))
824            .try_map(|_row: sqlx::sqlite::SqliteRow| Ok::<_, sqlx::Error>(42_i64));
825        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
826    }
827
828    #[test]
829    fn map_compose_after_annotations() {
830        // Multi-map composition: annotations survive across two `.map()` calls.
831        let q = sqlx::query::<Sqlite>("SELECT 1")
832            .with_annotations(QueryAnnotations::new().operation("SELECT"))
833            .map(|_row: sqlx::sqlite::SqliteRow| 1_i64)
834            .map(|n| n + 1);
835        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
836    }
837
838    #[test]
839    fn map_then_with_operation_replaces_via_wrapper() {
840        // `with_operation` shorthand on the new `AnnotatedQuery<Map<...>>` wrapper.
841        let q = sqlx::query::<Sqlite>("SELECT 1")
842            .map(|_row: sqlx::sqlite::SqliteRow| 1_i64)
843            .with_annotations(QueryAnnotations::new().query_summary("legacy"))
844            .with_operation("SELECT", "users");
845        assert_eq!(q.annotations.operation.as_deref(), Some("SELECT"));
846        assert_eq!(q.annotations.collection.as_deref(), Some("users"));
847        assert!(q.annotations.query_summary.is_none());
848    }
849
850    #[test]
851    fn debug_impl_for_annotated_map_includes_annotations() {
852        // The manual `Debug` impl on the new `Map` wrapper still prints annotations.
853        let q = sqlx::query::<Sqlite>("SELECT 1")
854            .map(|_row: sqlx::sqlite::SqliteRow| 1_i64)
855            .with_annotations(QueryAnnotations::new().operation("DEBUG_MAP"));
856        let debug = format!("{q:?}");
857        assert!(debug.contains("AnnotatedQuery"));
858        assert!(debug.contains("DEBUG_MAP"));
859    }
860}
sqlx_otel/query_ext.rs

sqlx_otel/
query_ext.rs
