wrap timed streams and iterators in tracing::Spans; follow Opentelemetry conventions

Cargo.toml

@@ -142,6 +142,7 @@ mac_address = "1.1.5"
 rust_decimal = { version = "1.26.1", default-features = false, features = ["std"] }
 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 }

sqlx-core/src/any/database.rs

@@ -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/options.rs

@@ -62,4 +62,9 @@ impl ConnectOptions for AnyConnectOptions {
         self.log_settings.slow_statements_duration = duration;
         self
     }
+
+    fn set_span_level(mut self, level: LevelFilter) -> Self {
+        self.log_settings.span_level = level;
+        self
+    }
 }

sqlx-core/src/connection.rs

@@ -161,6 +161,7 @@ pub struct LogSettings {
     pub statements_level: LevelFilter,
     pub slow_statements_level: LevelFilter,
     pub slow_statements_duration: Duration,
+    pub span_level: LevelFilter,
 }
 
 impl Default for LogSettings {
@@ -169,6 +170,7 @@ impl Default for LogSettings {
             statements_level: LevelFilter::Debug,
             slow_statements_level: LevelFilter::Warn,
             slow_statements_duration: Duration::from_secs(1),
+            span_level: LevelFilter::Info,
         }
     }
 }
@@ -181,6 +183,10 @@ impl LogSettings {
         self.slow_statements_level = level;
         self.slow_statements_duration = duration;
     }
+
+    pub fn set_span_level(&mut self, level: LevelFilter) {
+        self.span_level = level;
+    }
 }
 
 pub trait ConnectOptions: 'static + Send + Sync + FromStr<Err = Error> + Debug + Clone {
@@ -234,5 +240,8 @@ pub trait ConnectOptions: 'static + Send + Sync + FromStr<Err = Error> + Debug +
     fn disable_statement_logging(self) -> Self {
         self.log_statements(LevelFilter::Off)
             .log_slow_statements(LevelFilter::Off, Duration::default())
+            .set_span_level(LevelFilter::Off)
     }
+
+    fn set_span_level(self, level: LevelFilter) -> Self;
 }

sqlx-core/src/database.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/logger.rs

@@ -38,6 +38,22 @@ 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,19 +84,47 @@ 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 {
+        let span = Self::new_tracing_span(db_system, &settings);
+
         Self {
             sql,
             rows_returned: 0,
             rows_affected: 0,
             start: Instant::now(),
             settings,
+            span,
         }
     }
 
+    fn new_tracing_span(db_system: &str, settings: &LogSettings) -> tracing::Span {
+        // only create a usable span if the span level is set (i.e. not 'OFF') and
+        // the filter would output it
+        if let Some(level) = private_level_filter_to_trace_level(settings.span_level) {
+            if private_tracing_dynamic_enabled!(target: "sqlx::query", level) {
+                use tracing::field::Empty;
+
+                return 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,
+                );
+            }
+        }
+
+        // No-op span preferred over Option<Span>
+        tracing::Span::none()
+    }
+
     pub fn increment_rows_returned(&mut self) {
         self.rows_returned += 1;
     }
@@ -100,27 +144,50 @@ impl<'q> QueryLogger<'q> {
             self.settings.statements_level
         };
 
-        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
-            let log_is_enabled = log::log_enabled!(target: "sqlx::query", log_level)
-                || private_tracing_dynamic_enabled!(target: "sqlx::query", tracing_level);
-            if log_is_enabled {
-                let mut summary = parse_query_summary(self.sql);
-
-                let sql = if summary != self.sql {
-                    summary.push_str(" …");
-                    format!(
-                        "\n\n{}\n",
-                        sqlformat::format(
-                            self.sql,
-                            &sqlformat::QueryParams::None,
-                            sqlformat::FormatOptions::default()
-                        )
+        let span_level = private_level_filter_to_trace_level(self.settings.span_level);
+        let log_levels = private_level_filter_to_levels(lvl);
+
+        let span_is_enabled = span_level
+            .map(|level| private_tracing_dynamic_enabled!(target: "sqlx::query", level))
+            .unwrap_or(false);
+
+        let log_is_enabled = log_levels
+            .map(|(tracing_level, log_level)| {
+                // 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
+                log::log_enabled!(target: "sqlx::query", log_level)
+                    || private_tracing_dynamic_enabled!(target: "sqlx::query", tracing_level)
+            })
+            .unwrap_or(false);
+
+        // only do these potentially expensive operations if the span or log will record them
+        if span_is_enabled || log_is_enabled {
+            let mut summary = parse_query_summary(self.sql);
+
+            let sql = if summary != self.sql {
+                summary.push_str(" …");
+                format!(
+                    "\n\n{}\n",
+                    sqlformat::format(
+                        self.sql,
+                        &sqlformat::QueryParams::None,
+                        sqlformat::FormatOptions::default()
                     )
-                } else {
-                    String::new()
-                };
+                )
+            } else {
+                String::new()
+            };
+
+            // in the case where span_is_enabled is false, these will no-op
+            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 log_is_enabled {
+                let (tracing_level, _) = log_levels.unwrap(); // previously checked to be some
 
                 if was_slow {
                     private_tracing_dynamic_event!(

sqlx-core/src/pool/inner.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-mysql/Cargo.toml

@@ -68,6 +68,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/src/connection/executor.rs

@@ -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);
@@ -240,7 +243,7 @@ impl MySqlConnection {
                     r#yield!(v);
                 }
             }
-        }))
+        }).instrument(span_handle))
     }
 }
 

sqlx-mysql/src/database.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/options/connect.rs

@@ -94,4 +94,9 @@ impl ConnectOptions for MySqlConnectOptions {
         self.log_settings.log_slow_statements(level, duration);
         self
     }
+
+    fn set_span_level(mut self, level: LevelFilter) -> Self {
+        self.log_settings.span_level = level;
+        self
+    }
 }

sqlx-postgres/Cargo.toml

@@ -31,6 +31,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/src/connection/executor.rs

@@ -17,7 +17,9 @@ use futures_core::stream::BoxStream;
 use futures_core::Stream;
 use futures_util::{pin_mut, TryStreamExt};
 use sqlx_core::arguments::Arguments;
+use sqlx_core::database::Database;
 use sqlx_core::Either;
+use tracing_futures::Instrument;
 use std::{borrow::Cow, sync::Arc};
 
 async fn prepare(
@@ -195,7 +197,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?;
@@ -362,7 +365,7 @@ impl PgConnection {
             }
 
             Ok(())
-        })
+        }.instrument(span_handle))
     }
 }