feat!: annotate SQL statements with OTel trace context via AtomicOperation

Cargo.toml

@@ -25,6 +25,9 @@ instrument = ["es-entity-macros/instrument", "dep:tracing"]
 [dependencies]
 es-entity-macros = { workspace = true }
 
+async-stream = { workspace = true }
+futures-core = { workspace = true }
+futures-util = { workspace = true }
 base64 = { workspace = true, optional = true }
 sqlx = { workspace = true }
 thiserror = { workspace = true }
@@ -53,6 +56,7 @@ tokio = { workspace = true }
 anyhow = { workspace = true }
 async-trait = { workspace = true }
 futures = { workspace = true }
+tracing-subscriber = { workspace = true }
 
 [workspace]
 resolver = "2"
@@ -65,6 +69,7 @@ members = [
 es-entity-macros = { path = "es-entity-macros", version = "0.11.2-dev" }
 
 anyhow = "1.0"
+async-stream = "0.3"
 async-graphql = { version = "8.0.0-rc.5", default-features = false }
 async-trait = "0.1"
 base64 = { version = "0.22" }
@@ -81,8 +86,11 @@ uuid = { version = "1.23", features = ["serde", "v7"] }
 im = { version = "15.1", features = ["serde"] }
 pin-project = "1.1"
 tracing = { version = "0.1.41", default-features = false }
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 tracing-opentelemetry = { version = "0.33.0", default-features = false }
 opentelemetry = { version = "0.32.0", default-features = false }
 opentelemetry_sdk = { version = "0.32.0", features = ["rt-tokio"] }
 futures = "0.3"
+futures-core = "0.3"
+futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
 parking_lot = "0.12"

src/lib.rs

@@ -36,6 +36,7 @@ pub mod one_time_executor;
 pub mod operation;
 pub mod pagination;
 pub mod query;
+pub mod sql_commenter;
 pub mod traits;
 
 pub mod prelude {

src/one_time_executor.rs

@@ -1,6 +1,37 @@
 //! Type-safe wrapper to ensure one database operation per executor.
+//!
+//! [`OneTimeExecutor`] also implements [`sqlx::Executor`]: every statement
+//! executed through it is annotated with the active span's `traceparent` as a
+//! trailing SQL comment (see [`crate::sql_commenter`]). This is what allows
+//! statements observed server-side (`pg_stat_activity`, Postgres logs) to be
+//! matched to distributed traces.
+//!
+//! The annotation rewrites the statement *text* only; bind arguments and row
+//! mapping flow through unchanged, so it is transparent to `sqlx::query!`
+//! macro-generated queries as well as dynamically built ones.
+//!
+//! # Trade-off: annotated statements bypass the prepared statement cache
+//!
+//! The trace context makes each annotated statement's text unique, so it can
+//! never match sqlx's per-connection prepared statement cache. Annotated
+//! statements are therefore executed with `persistent(false)` (the unnamed
+//! statement) — bypassing the cache rather than thrashing it with single-use
+//! entries. The cost is a server-side parse + plan per annotated execution.
+//! Annotation only happens for *sampled* spans (see
+//! [`crate::sql_commenter::current_traceparent`]), so un-sampled traffic keeps
+//! full prepared-statement reuse.
+//!
+//! When there is no sampled span context the original query is passed through
+//! untouched and the executor adds no overhead.
 
-use crate::{db, operation::AtomicOperation};
+use async_stream::try_stream;
+use futures_core::stream::BoxStream;
+use futures_util::{TryStreamExt, future::BoxFuture};
+use sqlx::{Database, Describe, Error, Execute, Executor};
+
+use std::borrow::Cow;
+
+use crate::{db, operation::AtomicOperation, sql_commenter};
 
 /// A struct that owns an [`sqlx::Executor`].
 ///
@@ -21,9 +52,10 @@ use crate::{db, operation::AtomicOperation};
 ///     Ok(())
 /// }
 /// ```
+#[derive(Debug)]
 pub struct OneTimeExecutor<'c, E>
 where
-    E: sqlx::Executor<'c, Database = db::Db>,
+    E: sqlx::Executor<'c, Database = db::Db> + 'c,
 {
     now: Option<chrono::DateTime<chrono::Utc>>,
     executor: E,
@@ -32,9 +64,9 @@ where
 
 impl<'c, E> OneTimeExecutor<'c, E>
 where
-    E: sqlx::Executor<'c, Database = db::Db>,
+    E: sqlx::Executor<'c, Database = db::Db> + 'c,
 {
-    fn new(executor: E, now: Option<chrono::DateTime<chrono::Utc>>) -> Self {
+    pub(crate) fn new(executor: E, now: Option<chrono::DateTime<chrono::Utc>>) -> Self {
         OneTimeExecutor {
             executor,
             now,
@@ -56,7 +88,7 @@ where
         O: Send + Unpin,
         A: 'q + Send + sqlx::IntoArguments<'q, db::Db>,
     {
-        query.fetch_one(self.executor).await
+        query.fetch_one(self).await
     }
 
     /// Proxy call to `query.fetch_all` but guarantees the inner executor will only be used once.
@@ -69,7 +101,7 @@ where
         O: Send + Unpin,
         A: 'q + Send + sqlx::IntoArguments<'q, sqlx::Postgres>,
     {
-        query.fetch_all(self.executor).await
+        query.fetch_all(self).await
     }
 
     /// Proxy call to `query.fetch_optional` but guarantees the inner executor will only be used once.
@@ -82,7 +114,99 @@ where
         O: Send + Unpin,
         A: 'q + Send + sqlx::IntoArguments<'q, sqlx::Postgres>,
     {
-        query.fetch_optional(self.executor).await
+        query.fetch_optional(self).await
+    }
+}
+
+/// Returns the annotated statement text, or `None` when there is no sampled
+/// span context (in which case the query should be delegated untouched).
+fn annotated_sql<'q>(query: &impl Execute<'q, db::Db>) -> Option<String> {
+    match sql_commenter::annotate_sql(query.sql()) {
+        Cow::Borrowed(_) => None,
+        Cow::Owned(sql) => Some(sql),
+    }
+}
+
+/// Rebuilds a query with annotated SQL, moving out its bind arguments.
+///
+/// The trace context makes each statement's text unique, so the returned query
+/// is marked non-persistent: it can never hit sqlx's per-connection prepared
+/// statement cache, and caching it would evict useful entries. The trade-off
+/// is a server-side parse + plan per annotated execution.
+fn rebuild_annotated<'q>(
+    annotated: &str,
+    mut query: impl Execute<'q, db::Db>,
+) -> Result<sqlx::query::Query<'_, db::Db, sqlx::postgres::PgArguments>, Error> {
+    let args = query
+        .take_arguments()
+        .map_err(Error::Encode)?
+        .unwrap_or_default();
+    Ok(sqlx::query_with::<db::Db, _>(annotated, args).persistent(false))
+}
+
+impl<'c, E> Executor<'c> for OneTimeExecutor<'c, E>
+where
+    E: Executor<'c, Database = db::Db> + 'c,
+{
+    type Database = db::Db;
+
+    fn fetch_many<'e, 'q: 'e, Q>(
+        self,
+        query: Q,
+    ) -> BoxStream<'e, Result<sqlx::Either<<db::Db as Database>::QueryResult, db::Row>, Error>>
+    where
+        'c: 'e,
+        Q: 'q + Execute<'q, db::Db>,
+    {
+        let Some(annotated) = annotated_sql(&query) else {
+            return self.executor.fetch_many(query);
+        };
+        Box::pin(try_stream! {
+            let q = rebuild_annotated(annotated.as_str(), query)?;
+            let mut stream = self.executor.fetch_many(q);
+            while let Some(step) = stream.try_next().await? {
+                yield step;
+            }
+        })
+    }
+
+    fn fetch_optional<'e, 'q: 'e, Q>(
+        self,
+        query: Q,
+    ) -> BoxFuture<'e, Result<Option<db::Row>, Error>>
+    where
+        'c: 'e,
+        Q: 'q + Execute<'q, db::Db>,
+    {
+        let Some(annotated) = annotated_sql(&query) else {
+            return self.executor.fetch_optional(query);
+        };
+        Box::pin(async move {
+            let q = rebuild_annotated(annotated.as_str(), query)?;
+            self.executor.fetch_optional(q).await
+        })
+    }
+
+    fn prepare_with<'e, 'q: 'e>(
+        self,
+        sql: &'q str,
+        parameters: &'e [<db::Db as Database>::TypeInfo],
+    ) -> BoxFuture<'e, Result<<db::Db as Database>::Statement<'q>, Error>>
+    where
+        'c: 'e,
+    {
+        // A prepared Statement<'q> may borrow `sql`; a locally allocated
+        // annotated string could not satisfy the 'q lifetime, so preparation
+        // is delegated unannotated.
+        self.executor.prepare_with(sql, parameters)
+    }
+
+    fn describe<'e, 'q: 'e>(self, sql: &'q str) -> BoxFuture<'e, Result<Describe<db::Db>, Error>>
+    where
+        'c: 'e,
+    {
+        // Not an execution path; delegated unannotated.
+        self.executor.describe(sql)
     }
 }
 
@@ -149,6 +273,6 @@ where
         Self: 'c,
     {
         let now = self.maybe_now();
-        OneTimeExecutor::new(self.as_executor(), now)
+        OneTimeExecutor::new(self.connection(), now)
     }
 }

src/operation/hooks.rs

@@ -210,7 +210,7 @@ impl<'c> HookOperation<'c> {
     fn new(op: &'c mut impl AtomicOperation) -> Self {
         Self {
             now: op.maybe_now(),
-            conn: op.as_executor(),
+            conn: op.connection(),
         }
     }
 }
@@ -220,7 +220,7 @@ impl<'c> AtomicOperation for HookOperation<'c> {
         self.now
     }
 
-    fn as_executor(&mut self) -> &mut db::Connection {
+    fn connection(&mut self) -> &mut db::Connection {
         self.conn
     }
 }