From 7fb8306acb61890a253e22fe86cf7b78d1943452 Mon Sep 17 00:00:00 2001 From: Borislav Borisov Date: Sun, 3 May 2026 21:33:22 +0100 Subject: [PATCH] fix: Make query-side AnnotatedQuery forwarders Send-clean MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fetch_*/execute methods on AnnotatedQuery were declared as `async fn`, which generated a coroutine state machine on this crate's side. Inside that coroutine, rustc's auto-trait Send inference had to discharge the HRTB `for<'a> &'a mut DB::Connection: sqlx::Executor<'a, ...>` carried by the IntoAnnotatedExecutor impls – and could not, breaking query-side annotation inside tokio::spawn, axum handlers, and tower::Service-bounded futures. Convert the five forwarders (fetch_optional, fetch_all, fetch_one, execute, execute_many) from `async fn -> T` to `fn -> impl Future + Send + 'e`. The body forwards sqlx's underlying future directly with no `.await`, so no coroutine forms at this layer; Send delegates to sqlx's own already-Send future. The HRTB on the IntoAnnotatedExecutor impls and on impl_executor! is unchanged, and the executor-side surface is unaffected. --- CHANGELOG.md | 4 + src/query_ext.rs | 74 +++++++++++--- tests/common.rs | 244 ++++++++++++++++++++++++++++++++++++++++++---- tests/mysql.rs | 6 ++ tests/postgres.rs | 6 ++ tests/sqlite.rs | 6 ++ 6 files changed, 308 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 65a6580..d12f05c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ## [Unreleased] +### Fixed + +- Query-side `with_annotations` / `with_operation` now compile inside `Send`-required async contexts (axum handlers, `tokio::spawn`, `tower::Service`-bounded futures). The `AnnotatedQuery::fetch_*` / `execute` forwarders were rewritten from `async fn` into `fn -> impl Future + Send + 'e`, so the HRTB carried by the internal `IntoAnnotatedExecutor` impls no longer leaks into auto-trait inference of an opaque coroutine. Span output is byte-identical to the executor-side surface; metrics emission is unchanged by construction because both surfaces continue to funnel through the same `Annotated<'_, Pool>` `Executor` impl ([#NN](https://github.com/chmodas/sqlx-otel/pull/31)). + ## [0.2.0] – 2026-04-28 ### Added diff --git a/src/query_ext.rs b/src/query_ext.rs index 4933964..9eda556 100644 --- a/src/query_ext.rs +++ b/src/query_ext.rs @@ -130,6 +130,33 @@ pub trait QueryAnnotateExt: sealed::Sealed + Sized { /// /// Equivalent to /// `self.with_annotations(QueryAnnotations::new().operation(op).collection(coll))`. + /// + /// The returned future is `Send` and may be used inside a `Send`-required async context + /// (e.g. `tokio::spawn`, axum handlers, `tower::Service`-bounded futures). The example + /// below intentionally swallows the inner `sqlx::Error` via `.ok().flatten()` – the + /// point is to exercise the `Send` contract on the spawned future at compile time, not + /// to demonstrate error handling. A non-`Send` future would fail to compile here, since + /// `tokio::spawn` requires `Send`. + /// + /// ```no_run + /// # #[cfg(feature = "sqlite")] + /// # async fn _doc() -> Result<(), Box> { + /// # use sqlx_otel::PoolBuilder; + /// use sqlx_otel::QueryAnnotateExt as _; + /// # let pool: sqlx_otel::Pool = + /// # PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build(); + /// + /// tokio::spawn(async move { + /// let _: Option<(i32,)> = sqlx::query_as("SELECT 1") + /// .with_operation("SELECT", "users") + /// .fetch_optional(&pool) + /// .await + /// .ok() + /// .flatten(); + /// }) + /// .await?; + /// # Ok(()) } + /// ``` fn with_operation( self, operation: impl Into, @@ -215,9 +242,17 @@ impl std::fmt::Debug for AnnotatedQuery { /// `impl_executor!` already supports (`&Pool`, `&mut PoolConnection`, `&mut Transaction`), /// each producing the matching [`Annotated`] / [`AnnotatedMut`] wrapper. /// -/// The HRTB `for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>` lives on -/// each individual impl rather than on the trait, to avoid trait-resolution recursion when -/// the user's executor type is itself constructed via the same HRTB. +/// The HRTB `for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>` lives on each +/// individual impl rather than on the trait, to avoid trait-resolution recursion when the +/// user's executor type is itself constructed via the same HRTB. This bound used to leak into +/// the `Send` auto-trait inference of `AnnotatedQuery::fetch_*` / `execute` when those methods +/// were `async fn` – the resulting opaque coroutine carried a region-quantified obligation +/// that rustc could not always discharge, which broke query-side annotation inside contexts +/// requiring `Send` (`tokio::spawn`, axum handlers). The fix is in the *callers* of this +/// trait, not the trait itself: the forwarders on `AnnotatedQuery` are written as +/// `fn(...) -> impl Future + Send + 'e` returning `SQLx`'s own future directly, so no +/// coroutine forms on this crate's side and the `Send` check delegates to the underlying +/// `SQLx` future, which is already `Send`-clean for concrete backends. /// /// Users do not call this trait directly – they pass `&pool`, `&mut conn`, or `&mut tx` to /// [`AnnotatedQuery::execute`] / `fetch*` and the trait dispatches internally. The trait is @@ -360,7 +395,10 @@ macro_rules! impl_annotated_query_fetch_forwarders { /// /// Returns any [`sqlx::Error`] surfaced by the underlying driver, including row /// decoding errors. - pub async fn fetch_all<'e, E>(self, executor: E) -> Result, sqlx::Error> + pub fn fetch_all<'e, E>( + self, + executor: E, + ) -> impl 'e + Send + std::future::Future, sqlx::Error>> where 'q: 'e, A: 'e, @@ -368,7 +406,7 @@ macro_rules! impl_annotated_query_fetch_forwarders { E: 'e + IntoAnnotatedExecutor<'e, DB>, { let wrapper = executor.into_annotated(self.annotations); - self.inner.fetch_all(wrapper).await + self.inner.fetch_all(wrapper) } /// Return exactly one row, erroring if none or more than one. @@ -377,7 +415,10 @@ macro_rules! impl_annotated_query_fetch_forwarders { /// /// Returns [`sqlx::Error::RowNotFound`] when the result set is empty, or any other /// [`sqlx::Error`] surfaced by the underlying driver. - pub async fn fetch_one<'e, E>(self, executor: E) -> Result<$row, sqlx::Error> + pub fn fetch_one<'e, E>( + self, + executor: E, + ) -> impl 'e + Send + std::future::Future> where 'q: 'e, A: 'e, @@ -385,7 +426,7 @@ macro_rules! impl_annotated_query_fetch_forwarders { E: 'e + IntoAnnotatedExecutor<'e, DB>, { let wrapper = executor.into_annotated(self.annotations); - self.inner.fetch_one(wrapper).await + self.inner.fetch_one(wrapper) } /// Return at most one row. @@ -393,10 +434,10 @@ macro_rules! impl_annotated_query_fetch_forwarders { /// # Errors /// /// Returns any [`sqlx::Error`] surfaced by the underlying driver. - pub async fn fetch_optional<'e, E>( + pub fn fetch_optional<'e, E>( self, executor: E, - ) -> Result, sqlx::Error> + ) -> impl 'e + Send + std::future::Future, sqlx::Error>> where 'q: 'e, A: 'e, @@ -404,7 +445,7 @@ macro_rules! impl_annotated_query_fetch_forwarders { E: 'e + IntoAnnotatedExecutor<'e, DB>, { let wrapper = executor.into_annotated(self.annotations); - self.inner.fetch_optional(wrapper).await + self.inner.fetch_optional(wrapper) } }; } @@ -442,14 +483,17 @@ where /// # Errors /// /// Returns any [`sqlx::Error`] surfaced by the underlying driver. - pub async fn execute<'e, E>(self, executor: E) -> Result + pub fn execute<'e, E>( + self, + executor: E, + ) -> impl 'e + Send + std::future::Future> where 'q: 'e, A: 'e, E: 'e + IntoAnnotatedExecutor<'e, DB>, { let wrapper = executor.into_annotated(self.annotations); - self.inner.execute(wrapper).await + self.inner.execute(wrapper) } /// Execute multiple statements separated by `;` and return their results as a stream. @@ -458,17 +502,17 @@ where /// existing executor-side surface. Only `Query` exposes this method – `QueryAs`, /// `QueryScalar`, and `Map` have no `execute_many` upstream. #[allow(deprecated)] - pub async fn execute_many<'e, E>( + pub fn execute_many<'e, E>( self, executor: E, - ) -> BoxStream<'e, Result> + ) -> impl 'e + Send + std::future::Future>> where 'q: 'e, A: 'e, E: 'e + IntoAnnotatedExecutor<'e, DB>, { let wrapper = executor.into_annotated(self.annotations); - self.inner.execute_many(wrapper).await + self.inner.execute_many(wrapper) } /// Map each row to another type. Mirrors [`sqlx::query::Query::map`] and carries the diff --git a/tests/common.rs b/tests/common.rs index 28cc982..eb5345a 100644 --- a/tests/common.rs +++ b/tests/common.rs @@ -323,6 +323,50 @@ pub fn assert_one_annotated_span(tel: &TestTelemetry, dialect: &Dialect) { assert_annotated_span(&spans[0], dialect); } +/// Compare two single-span exporter snapshots and assert that the executor-side and +/// query-side annotation surfaces emitted byte-identical span data on the dimensions +/// users observe (name, kind, and the full annotation/connection attribute set). The +/// `case` argument is included in failure messages so the cause is obvious when one of +/// several builder-family pairs in a parity test fails. +pub fn assert_span_parity(case: &str, exec_spans: &[SpanData], query_spans: &[SpanData]) { + assert_eq!( + exec_spans.len(), + 1, + "{case}: executor-side emitted {} spans, expected 1", + exec_spans.len() + ); + assert_eq!( + query_spans.len(), + 1, + "{case}: query-side emitted {} spans, expected 1", + query_spans.len() + ); + let exec = &exec_spans[0]; + let query = &query_spans[0]; + + assert_eq!(exec.name, query.name, "{case}: span name differs"); + assert_eq!(exec.span_kind, query.span_kind, "{case}: span kind differs"); + for key in &[ + "db.system.name", + "db.operation.name", + "db.collection.name", + "db.query.text", + "db.query.summary", + "db.namespace", + "db.stored_procedure.name", + "server.address", + "server.port", + "db.response.affected_rows", + "db.response.returned_rows", + ] { + assert_eq!( + attr(exec, key), + attr(query, key), + "{case}: attribute `{key}` differs across executor-side vs query-side", + ); + } +} + // --------------------------------------------------------------------------- // Parameterised test bodies (macro_rules) // --------------------------------------------------------------------------- @@ -2213,13 +2257,28 @@ macro_rules! test_query_as_fetch_all_with_annotations_via_pool { .unwrap(); assert_eq!(rows.len(), 2); + let pool_clone = pool.clone(); + let rows: Vec<(i32,)> = tokio::spawn(async move { + sqlx::query_as("SELECT 1 UNION ALL SELECT 2") + .with_annotations($crate::common::test_annotations()) + .fetch_all(&pool_clone) + .await + .unwrap() + }) + .await + .unwrap(); + assert_eq!(rows.len(), 2); + let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + } }}; } -/// `sqlx::query_as(...).with_annotations(...).fetch_one(&pool)`. +/// `sqlx::query_as(...).with_annotations(...).fetch_one(&pool)`. Runs inline and inside +/// `tokio::spawn` to exercise the `Send`-required path. #[macro_export] macro_rules! test_query_as_fetch_one_with_annotations_via_pool { ($pool_factory:expr, $dialect:expr) => {{ @@ -2234,13 +2293,28 @@ macro_rules! test_query_as_fetch_one_with_annotations_via_pool { .unwrap(); assert_eq!(row.0, 7); + let pool_clone = pool.clone(); + let row: (i32,) = tokio::spawn(async move { + sqlx::query_as("SELECT 7") + .with_annotations($crate::common::test_annotations()) + .fetch_one(&pool_clone) + .await + .unwrap() + }) + .await + .unwrap(); + assert_eq!(row.0, 7); + let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + } }}; } /// `sqlx::query_as(...).with_annotations(...).fetch_optional(&pool)` returning none. +/// Runs inline and inside `tokio::spawn` to exercise the `Send`-required path. #[macro_export] macro_rules! test_query_as_fetch_optional_with_annotations_via_pool { ($pool_factory:expr, $dialect:expr) => {{ @@ -2255,9 +2329,23 @@ macro_rules! test_query_as_fetch_optional_with_annotations_via_pool { .unwrap(); assert!(row.is_none()); + let pool_clone = pool.clone(); + let row: Option<(i32,)> = tokio::spawn(async move { + sqlx::query_as("SELECT 1 WHERE 1 = 0") + .with_annotations($crate::common::test_annotations()) + .fetch_optional(&pool_clone) + .await + .unwrap() + }) + .await + .unwrap(); + assert!(row.is_none()); + let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + } }}; } @@ -3147,7 +3235,11 @@ macro_rules! test_query_annotations_first_then_bind_via_pool { } /// Annotated `execute` against the wrapped pool. Uses `SELECT 1` so the test is -/// portable; the executor records `affected_rows` regardless of statement kind. +/// portable; the executor records `affected_rows` regardless of statement kind. The +/// macro runs the call twice – once inline (documents the simplest usage) and once +/// inside `tokio::spawn` (compile-time proof that the returned future is `Send`, the +/// contract that broke under v0.2.0's `async fn` shape and was restored by converting +/// the forwarder to `fn -> impl Future + Send + 'e`). #[macro_export] macro_rules! test_query_execute_with_annotations_via_pool { ($pool_factory:expr, $dialect:expr) => {{ @@ -3161,14 +3253,29 @@ macro_rules! test_query_execute_with_annotations_via_pool { .await .unwrap(); + let pool_clone = pool.clone(); + tokio::spawn(async move { + sqlx::query("SELECT 1") + .with_annotations($crate::common::test_annotations()) + .execute(&pool_clone) + .await + .unwrap(); + }) + .await + .unwrap(); + let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); - assert!($crate::common::attr(&spans[0], "db.response.affected_rows").is_some()); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + assert!($crate::common::attr(span, "db.response.affected_rows").is_some()); + } }}; } -/// Annotated `execute` against `&mut PoolConnection`. +/// Annotated `execute` against `&mut PoolConnection`. Runs inline and inside +/// `tokio::spawn` to exercise the `Send`-required path over the `&mut PoolConnection` +/// borrow. #[macro_export] macro_rules! test_query_execute_with_annotations_via_connection { ($pool_factory:expr, $dialect:expr) => {{ @@ -3182,14 +3289,31 @@ macro_rules! test_query_execute_with_annotations_via_connection { .execute(&mut conn) .await .unwrap(); + drop(conn); + + let pool_clone = pool.clone(); + tokio::spawn(async move { + let mut conn = pool_clone.acquire().await.unwrap(); + sqlx::query("SELECT 1") + .with_annotations($crate::common::test_annotations()) + .execute(&mut conn) + .await + .unwrap(); + }) + .await + .unwrap(); let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + } }}; } -/// Annotated `execute` against `&mut Transaction<'_, DB>`. +/// Annotated `execute` against `&mut Transaction<'_, DB>`. Runs inline and inside +/// `tokio::spawn` to exercise the `Send`-required path over the `&mut Transaction` +/// borrow. #[macro_export] macro_rules! test_query_execute_with_annotations_via_transaction { ($pool_factory:expr, $dialect:expr) => {{ @@ -3205,9 +3329,24 @@ macro_rules! test_query_execute_with_annotations_via_transaction { .unwrap(); tx.commit().await.unwrap(); + let pool_clone = pool.clone(); + tokio::spawn(async move { + let mut tx = pool_clone.begin().await.unwrap(); + sqlx::query("SELECT 1") + .with_annotations($crate::common::test_annotations()) + .execute(&mut tx) + .await + .unwrap(); + tx.commit().await.unwrap(); + }) + .await + .unwrap(); + let spans = tel.spans(); - assert_eq!(spans.len(), 1); - $crate::common::assert_annotated_span(&spans[0], &$dialect); + assert_eq!(spans.len(), 2); + for span in &spans { + $crate::common::assert_annotated_span(span, &$dialect); + } }}; } @@ -3257,3 +3396,74 @@ macro_rules! test_query_fetch_optional_with_annotations_via_pool { ); }}; } + +/// Parity assertion: executor-side and query-side annotation surfaces must produce +/// byte-identical span attributes (per [`crate::query_ext`]'s "Choosing between +/// executor-side and query-side" doc, the two surfaces are documented as semantically +/// equivalent). Both code paths funnel through the same `Annotated<'_, Pool>` +/// `Executor` impl; if a future refactor diverges them, this test fails. Exercises all +/// three builder families (`query`, `query_as`, `query_scalar`) so a regression in any +/// one of them is caught. +#[macro_export] +macro_rules! test_executor_side_query_side_parity_via_pool { + ($pool_factory:expr, $dialect:expr) => {{ + use sqlx_otel::QueryAnnotateExt as _; + let pool = $pool_factory; + let tel = $crate::common::TestTelemetry::install(); + + // --- query (execute) ------------------------------------------------- + let _ = sqlx::query("SELECT 1") + .execute(pool.with_operation("SELECT", "users")) + .await + .unwrap(); + let exec_spans = tel.spans(); + tel.reset(); + + let _ = sqlx::query("SELECT 1") + .with_operation("SELECT", "users") + .execute(&pool) + .await + .unwrap(); + let query_spans = tel.spans(); + tel.reset(); + + $crate::common::assert_span_parity("query::execute", &exec_spans, &query_spans); + + // --- query_as (fetch_optional) --------------------------------------- + let _: Option<(i32,)> = sqlx::query_as("SELECT 1") + .fetch_optional(pool.with_operation("SELECT", "users")) + .await + .unwrap(); + let exec_spans = tel.spans(); + tel.reset(); + + let _: Option<(i32,)> = sqlx::query_as("SELECT 1") + .with_operation("SELECT", "users") + .fetch_optional(&pool) + .await + .unwrap(); + let query_spans = tel.spans(); + tel.reset(); + + $crate::common::assert_span_parity("query_as::fetch_optional", &exec_spans, &query_spans); + + // --- query_scalar (fetch_one) ---------------------------------------- + let _: i32 = sqlx::query_scalar("SELECT 1") + .fetch_one(pool.with_operation("SELECT", "users")) + .await + .unwrap(); + let exec_spans = tel.spans(); + tel.reset(); + + let _: i32 = sqlx::query_scalar("SELECT 1") + .with_operation("SELECT", "users") + .fetch_one(&pool) + .await + .unwrap(); + let query_spans = tel.spans(); + + $crate::common::assert_span_parity("query_scalar::fetch_one", &exec_spans, &query_spans); + + let _ = $dialect; + }}; +} diff --git a/tests/mysql.rs b/tests/mysql.rs index e4cb0ac..c741822 100644 --- a/tests/mysql.rs +++ b/tests/mysql.rs @@ -1095,3 +1095,9 @@ async fn query_scalar_macro_fetch_all_with_annotations_via_pool() { common::assert_one_annotated_span(&tel, &common::MYSQL_DIALECT); } + +#[tokio::test] +#[serial] +async fn executor_side_query_side_parity_via_pool() { + test_executor_side_query_side_parity_via_pool!(test_pool().await, common::MYSQL_DIALECT); +} diff --git a/tests/postgres.rs b/tests/postgres.rs index d382d9d..e1a3a78 100644 --- a/tests/postgres.rs +++ b/tests/postgres.rs @@ -1115,3 +1115,9 @@ async fn query_scalar_macro_fetch_all_with_annotations_via_pool() { common::assert_one_annotated_span(&tel, &common::POSTGRES_DIALECT); } + +#[tokio::test] +#[serial] +async fn executor_side_query_side_parity_via_pool() { + test_executor_side_query_side_parity_via_pool!(test_pool().await, common::POSTGRES_DIALECT); +} diff --git a/tests/sqlite.rs b/tests/sqlite.rs index 6af5cfa..2139423 100644 --- a/tests/sqlite.rs +++ b/tests/sqlite.rs @@ -1051,3 +1051,9 @@ async fn query_span_inherits_parent_context_streaming() { "stream query span parent_span_id should equal the outer span's span_id", ); } + +#[tokio::test] +#[serial] +async fn executor_side_query_side_parity_via_pool() { + test_executor_side_query_side_parity_via_pool!(test_pool().await, common::SQLITE_DIALECT); +}