feat(tracing): add InstrumentedStream span combinator to sqlx-core

Review finding #1 (shared mechanism): tracing 0.1's `Instrumented` only implements `Future`, not `Stream`, so a `#[tracing::instrument]` on an async fn that returns a stream closes its span before the stream is ever polled. Add a small pin-projected `InstrumentedStream` adapter and `InstrumentStream` extension trait that enters a `Span` on every `poll_next`, letting a span cover the full lifetime of a streamed query. pin-project-lite is already in the dependency tree.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: grahamc

sqlx-core/Cargo.toml

@@ -88,6 +88,7 @@ futures-util = { version = "0.3.32", default-features = false, features = ["allo
 log = { version = "0.4.18", default-features = false }
 memchr = { version = "2.5.0", default-features = false }
 percent-encoding = "2.3.0"
+pin-project-lite = "0.2.16"
 serde = { version = "1.0.219", features = ["derive", "rc"], optional = true }
 serde_json = { version = "1.0.142", features = ["raw_value"], optional = true }
 toml = { version = "0.8.16", optional = true }

sqlx-core/src/instrument_stream.rs

@@ -0,0 +1,52 @@
+//! Attach a [`tracing::Span`] to a [`Stream`] so the span stays open for the
+//! whole stream rather than just its construction.
+
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures_core::Stream;
+use pin_project_lite::pin_project;
+use tracing::Span;
+
+pin_project! {
+    /// A [`Stream`] adapter that enters `span` for the duration of every
+    /// [`poll_next`](Stream::poll_next).
+    ///
+    /// A plain `#[tracing::instrument]` on an `async fn` that *returns* a stream
+    /// only keeps its span entered while the stream is being built; the span is
+    /// then closed before the caller ever polls the stream. Wrapping the
+    /// returned stream with this adapter instead keeps the span open across row
+    /// fetching, so e.g. a `sqlx::query` span measures the whole query rather
+    /// than just its setup.
+    pub struct InstrumentedStream<S> {
+        #[pin]
+        stream: S,
+        span: Span,
+    }
+}
+
+impl<S: Stream> Stream for InstrumentedStream<S> {
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.project();
+        let _entered = this.span.enter();
+        this.stream.poll_next(cx)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.stream.size_hint()
+    }
+}
+
+/// Extension trait for attaching a [`Span`] to a [`Stream`].
+pub trait InstrumentStream: Stream + Sized {
+    /// Wrap this stream so that `span` is entered every time it is polled.
+    ///
+    /// See [`InstrumentedStream`].
+    fn instrument_stream(self, span: Span) -> InstrumentedStream<Self> {
+        InstrumentedStream { stream: self, span }
+    }
+}
+
+impl<S: Stream> InstrumentStream for S {}

sqlx-core/src/lib.rs

@@ -67,6 +67,7 @@ pub mod executor;
 pub mod from_row;
 pub mod fs;
 pub mod io;
+pub mod instrument_stream;
 pub mod logger;
 pub mod net;
 pub mod query_as;