wrap timed streams and iterators in tracing::Spans

Author: zvkemp

Cargo.lock

@@ -138,6 +138,7 @@ mac_address = "1.1.5"
 rust_decimal = "1.26.1"
 time = { version = "0.3.36", features = ["formatting", "parsing", "macros"] }
 uuid = "1.1.2"
+tracing-futures = { version = "0.2.5", features = ["futures-03", "std-future"] }
 
 # Common utility crates
 dotenvy = { version = "0.15.0", default-features = false }

Cargo.toml

@@ -31,6 +31,7 @@ impl Database for Any {
     type Statement<'q> = AnyStatement<'q>;
 
     const NAME: &'static str = "Any";
+    const NAME_LOWERCASE: &'static str = "any";
 
     const URL_SCHEMES: &'static [&'static str] = &[];
 }

sqlx-core/src/any/database.rs

@@ -181,6 +181,10 @@ impl LogSettings {
         self.slow_statements_level = level;
         self.slow_statements_duration = duration;
     }
+
+    pub fn tracing_span_level(&self) -> LevelFilter {
+        std::cmp::max(self.slow_statements_level, self.statements_level)
+    }
 }
 
 pub trait ConnectOptions: 'static + Send + Sync + FromStr<Err = Error> + Debug + Clone {

sqlx-core/src/connection.rs

@@ -105,6 +105,7 @@ pub trait Database: 'static + Sized + Send + Debug {
 
     /// The display name for this database driver.
     const NAME: &'static str;
+    const NAME_LOWERCASE: &'static str;
 
     /// The schemes for database URLs that should match this driver.
     const URL_SCHEMES: &'static [&'static str];

sqlx-core/src/database.rs

@@ -38,6 +38,23 @@ macro_rules! private_tracing_dynamic_event {
     }};
 }
 
+#[doc(hidden)]
+#[macro_export]
+macro_rules! private_tracing_dynamic_span {
+    (target: $target:expr, $level:expr, $($args:tt)*) => {{
+        use ::tracing::Level;
+
+        match $level {
+            Level::ERROR => ::tracing::span!(target: $target, Level::ERROR, $($args)*),
+            Level::WARN => ::tracing::span!(target: $target, Level::WARN, $($args)*),
+            Level::INFO => ::tracing::span!(target: $target, Level::INFO, $($args)*),
+            Level::DEBUG => ::tracing::span!(target: $target, Level::DEBUG, $($args)*),
+            Level::TRACE => ::tracing::span!(target: $target, Level::TRACE, $($args)*),
+        }
+    }};
+}
+
+
 #[doc(hidden)]
 pub fn private_level_filter_to_levels(
     filter: log::LevelFilter,
@@ -68,16 +85,33 @@ pub struct QueryLogger<'q> {
     rows_affected: u64,
     start: Instant,
     settings: LogSettings,
+    pub span: tracing::Span,
 }
 
 impl<'q> QueryLogger<'q> {
-    pub fn new(sql: &'q str, settings: LogSettings) -> Self {
+    pub fn new(sql: &'q str, db_system: &'q str, settings: LogSettings) -> Self {
+        use tracing::field::Empty;
+
+        let level = private_level_filter_to_trace_level(settings.tracing_span_level()).unwrap_or(tracing::Level::TRACE);
+
+        let span = private_tracing_dynamic_span!(
+            target: "sqlx::query",
+            level, 
+            "sqlx::query",
+            summary = Empty,
+            db.statement = Empty,
+            db.system = db_system,
+            rows_affected = Empty,
+            rows_returned = Empty,
+        );
+
         Self {
             sql,
             rows_returned: 0,
             rows_affected: 0,
             start: Instant::now(),
             settings,
+            span,
         }
     }
 
@@ -100,6 +134,7 @@ impl<'q> QueryLogger<'q> {
             self.settings.statements_level
         };
 
+        // only create an event if the level filter is set
         if let Some((tracing_level, log_level)) = private_level_filter_to_levels(lvl) {
             // The enabled level could be set from either tracing world or log world, so check both
             // to see if logging should be enabled for our level
@@ -122,37 +157,30 @@ impl<'q> QueryLogger<'q> {
                     String::new()
                 };
 
-                if was_slow {
+                self.span.record("summary", summary);
+                self.span.record("db.statement", sql);
+                self.span.record("rows_affected", self.rows_affected);
+                self.span.record("rows_returned", self.rows_returned);
+
+                let _e = self.span.enter();
+                if was_slow { 
                     private_tracing_dynamic_event!(
                         target: "sqlx::query",
                         tracing_level,
-                        summary,
-                        db.statement = sql,
-                        rows_affected = self.rows_affected,
-                        rows_returned = self.rows_returned,
-                        // Human-friendly - includes units (usually ms). Also kept for backward compatibility
-                        ?elapsed,
-                        // Search friendly - numeric
-                        elapsed_secs = elapsed.as_secs_f64(),
-                        // When logging to JSON, one can trigger alerts from the presence of this field.
                         slow_threshold=?self.settings.slow_statements_duration,
                         // Make sure to use "slow" in the message as that's likely
                         // what people will grep for.
                         "slow statement: execution time exceeded alert threshold"
-                    );
+                    )
                 } else {
                     private_tracing_dynamic_event!(
-                        target: "sqlx::query",
-                        tracing_level,
-                        summary,
-                        db.statement = sql,
-                        rows_affected = self.rows_affected,
-                        rows_returned = self.rows_returned,
-                        // Human-friendly - includes units (usually ms). Also kept for backward compatibility
-                        ?elapsed,
-                        // Search friendly - numeric
-                        elapsed_secs = elapsed.as_secs_f64(),
-                    );
+                         target: "sqlx::query",
+                         tracing_level,
+                         // Human-friendly - includes units (usually ms). Also kept for backward compatibility
+                         ?elapsed,
+                         // Search friendly - numeric
+                         elapsed_secs = elapsed.as_secs_f64(),
+                    )
                 }
             }
         }

sqlx-core/src/logger.rs

@@ -4,7 +4,9 @@ use crate::connection::Connection;
 use crate::database::Database;
 use crate::error::Error;
 use crate::pool::{deadline_as_timeout, CloseEvent, Pool, PoolOptions};
+use crate::private_tracing_dynamic_span;
 use crossbeam_queue::ArrayQueue;
+use tracing::Instrument;
 
 use crate::sync::{AsyncSemaphore, AsyncSemaphoreReleaser};
 
@@ -245,6 +247,10 @@ impl<DB: Database> PoolInner<DB> {
         }
     }
 
+    pub fn tracing_span_level(&self) -> Level {
+        std::cmp::max(self.acquire_slow_level, self.acquire_time_level).unwrap_or(Level::TRACE)
+    }
+
     pub(super) async fn acquire(self: &Arc<Self>) -> Result<Floating<DB, Live<DB>>, Error> {
         if self.is_closed() {
             return Err(Error::PoolClosed);
@@ -253,6 +259,13 @@ impl<DB: Database> PoolInner<DB> {
         let acquire_started_at = Instant::now();
         let deadline = acquire_started_at + self.options.acquire_timeout;
 
+        let span = private_tracing_dynamic_span!(
+                target: "sqlx::pool::acquire",
+                self.tracing_span_level(),
+                "sqlx::pool::acquire",
+                db.system = DB::NAME_LOWERCASE,
+        );
+
         let acquired = crate::rt::timeout(
             self.options.acquire_timeout,
             async {
@@ -295,6 +308,7 @@ impl<DB: Database> PoolInner<DB> {
                 }
             }
         )
+            .instrument(span.clone())
             .await
             .map_err(|_| Error::PoolTimedOut)??;
 
@@ -304,6 +318,8 @@ impl<DB: Database> PoolInner<DB> {
             .acquire_slow_level
             .filter(|_| acquired_after > self.options.acquire_slow_threshold);
 
+        let _e = span.enter();
+
         if let Some(level) = acquire_slow_level {
             private_tracing_dynamic_event!(
                 target: "sqlx::pool::acquire",

sqlx-core/src/pool/inner.rs

@@ -61,6 +61,7 @@ smallvec = "1.7.0"
 stringprep = "0.1.2"
 thiserror = "1.0.35"
 tracing = { version = "0.1.37", features = ["log"] }
+tracing-futures.workspace = true
 whoami = { version = "1.2.1", default-features = false }
 
 serde = { version = "1.0.144", optional = true }

sqlx-mysql/Cargo.toml

@@ -22,6 +22,8 @@ use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
 use futures_core::Stream;
 use futures_util::{pin_mut, TryStreamExt};
+use sqlx_core::database::Database;
+use tracing_futures::Instrument;
 use std::{borrow::Cow, sync::Arc};
 
 impl MySqlConnection {
@@ -106,7 +108,8 @@ impl MySqlConnection {
         persistent: bool,
     ) -> Result<impl Stream<Item = Result<Either<MySqlQueryResult, MySqlRow>, Error>> + 'e, Error>
     {
-        let mut logger = QueryLogger::new(sql, self.inner.log_settings.clone());
+        let mut logger = QueryLogger::new(sql, MySql::NAME_LOWERCASE, self.inner.log_settings.clone());
+        let span_handle = logger.span.clone();
 
         self.inner.stream.wait_until_ready().await?;
         self.inner.stream.waiting.push_back(Waiting::Result);
@@ -238,7 +241,7 @@ impl MySqlConnection {
                     r#yield!(v);
                 }
             }
-        }))
+        }).instrument(span_handle))
     }
 }
 

sqlx-mysql/src/connection/executor.rs

@@ -31,6 +31,7 @@ impl Database for MySql {
     type Statement<'q> = MySqlStatement<'q>;
 
     const NAME: &'static str = "MySQL";
+    const NAME_LOWERCASE: &'static str = "mysql";
 
     const URL_SCHEMES: &'static [&'static str] = &["mysql", "mariadb"];
 }

sqlx-mysql/src/database.rs

@@ -25,6 +25,7 @@ futures-channel = { version = "0.3.19", default-features = false, features = ["s
 futures-core = { version = "0.3.19", default-features = false }
 futures-io = "0.3.24"
 futures-util = { version = "0.3.19", default-features = false, features = ["alloc", "sink", "io"] }
+tracing-futures.workspace = true
 
 # Cryptographic Primitives
 crc = "3.0.0"

sqlx-postgres/Cargo.toml

@@ -17,7 +17,9 @@ use futures_core::future::BoxFuture;
 use futures_core::stream::BoxStream;
 use futures_core::Stream;
 use futures_util::{pin_mut, TryStreamExt};
+use sqlx_core::database::Database;
 use sqlx_core::Either;
+use tracing_futures::Instrument;
 use std::{borrow::Cow, sync::Arc};
 
 async fn prepare(
@@ -198,7 +200,8 @@ impl PgConnection {
         persistent: bool,
         metadata_opt: Option<Arc<PgStatementMetadata>>,
     ) -> Result<impl Stream<Item = Result<Either<PgQueryResult, PgRow>, Error>> + 'e, Error> {
-        let mut logger = QueryLogger::new(query, self.log_settings.clone());
+        let mut logger = QueryLogger::new(query, Postgres::NAME_LOWERCASE, self.log_settings.clone());
+        let span_handle = logger.span.clone();
 
         // before we continue, wait until we are "ready" to accept more queries
         self.wait_until_ready().await?;
@@ -352,7 +355,7 @@ impl PgConnection {
             }
 
             Ok(())
-        })
+        }.instrument(span_handle))
     }
 }
 

sqlx-postgres/src/connection/executor.rs

@@ -6,13 +6,14 @@ use futures_channel::mpsc::UnboundedSender;
 use futures_util::SinkExt;
 use log::Level;
 use sqlx_core::bytes::{Buf, Bytes};
+use sqlx_core::database::Database;
 
 use crate::connection::tls::MaybeUpgradeTls;
 use crate::error::Error;
 use crate::io::{Decode, Encode};
 use crate::message::{Message, MessageFormat, Notice, Notification, ParameterStatus};
 use crate::net::{self, BufferedSocket, Socket};
-use crate::{PgConnectOptions, PgDatabaseError, PgSeverity};
+use crate::{PgConnectOptions, PgDatabaseError, PgSeverity, Postgres};
 
 // the stream is a separate type from the connection to uphold the invariant where an instantiated
 // [PgConnection] is a **valid** connection to postgres
@@ -162,6 +163,7 @@ impl PgStream {
                         sqlx_core::private_tracing_dynamic_event!(
                             target: "sqlx::postgres::notice",
                             tracing_level,
+                            db.system = Postgres::NAME_LOWERCASE,
                             message = notice.message()
                         );
                     }

sqlx-postgres/src/connection/stream.rs

@@ -33,6 +33,7 @@ impl Database for Postgres {
     type Statement<'q> = PgStatement<'q>;
 
     const NAME: &'static str = "PostgreSQL";
+    const NAME_LOWERCASE: &'static str = "postgresql";
 
     const URL_SCHEMES: &'static [&'static str] = &["postgres", "postgresql"];
 }