Add SQLx operation tracing (#5223)

* wip: vendor sqlx-tracing

* (compiles) standardize pg types used

* more standardization

* general log message improvements

* wip: improve sqlx-tracing architecture

* unify sqlx::Executor type

* wip: try fix sqlx tracing

* wip: sqlx-tracing compiles

* so close

* it compiles

* fix ci

packages/sqlx-tracing/src/lib.rs

@@ -0,0 +1,219 @@
+#![doc = include_str!("../README.md")]
+
+use std::sync::Arc;
+
+use derive_more::{Deref, DerefMut};
+use futures::future::BoxFuture;
+
+mod any_connection;
+mod connection;
+mod pool;
+pub(crate) mod span;
+mod transaction;
+
+pub use sqlx::Executor;
+
+#[cfg(feature = "postgres")]
+pub mod postgres;
+
+#[cfg(feature = "sqlite")]
+pub mod sqlite;
+
+/// Attributes describing the database connection and context.
+/// Used for span enrichment and attribute propagation.
+#[derive(Debug, Default)]
+struct Attributes {
+    name: Option<String>,
+    host: Option<String>,
+    port: Option<u16>,
+    database: Option<String>,
+}
+
+pub trait Database: sqlx::Database {
+    const SYSTEM: &'static str;
+
+    /// Defines the type of reference to a database connection, equivalent to
+    /// `&'c mut <Self as sqlx::Database>::Connection`.
+    ///
+    /// But we can't actually use the `sqlx::Database` named connection type,
+    /// since we can't statically prove that it implements `sqlx::Executor`.
+    /// Even if we unify the two types (see `any_connection.rs`), we can't use
+    /// connection refs as an executor. So we need this intermediate associated
+    /// type.
+    type ConnectionRef<'c>: sqlx::Executor<'c, Database = Self>;
+
+    /// Casts a `&'c mut Self::Connection` to a `Self::ConnectionRef<'c>`.
+    ///
+    /// This should just return `conn`.
+    fn cast_connection<'c>(
+        conn: &'c mut <Self as sqlx::Database>::Connection,
+    ) -> Self::ConnectionRef<'c>;
+}
+
+/// Builder for constructing a [`Pool`] with custom attributes.
+///
+/// Allows setting database name, host, port, and other identifying information
+/// for tracing purposes.
+#[derive(Debug)]
+pub struct PoolBuilder<DB: Database> {
+    pool: sqlx::Pool<DB>,
+    attributes: Attributes,
+}
+
+// this is required because `pool.connect_options().to_url_lossy()` panics with sqlite
+#[cfg(feature = "postgres")]
+impl From<sqlx::Pool<sqlx::Postgres>> for PoolBuilder<sqlx::Postgres> {
+    /// Create a new builder from an existing SQLx pool.
+    fn from(pool: sqlx::Pool<sqlx::Postgres>) -> Self {
+        use sqlx::ConnectOptions;
+
+        let url = pool.connect_options().to_url_lossy();
+        let attributes = Attributes {
+            name: None,
+            host: url.host_str().map(String::from),
+            port: url.port(),
+            database: url
+                .path_segments()
+                .and_then(|mut segments| segments.next().map(String::from)),
+        };
+        Self { pool, attributes }
+    }
+}
+
+// this is required because `pool.connect_options().to_url_lossy()` panics with sqlite
+#[cfg(feature = "sqlite")]
+impl From<sqlx::Pool<sqlx::Sqlite>> for PoolBuilder<sqlx::Sqlite> {
+    /// Create a new builder from an existing SQLx pool.
+    fn from(pool: sqlx::Pool<sqlx::Sqlite>) -> Self {
+        let attributes = Attributes {
+            name: None,
+            host: pool
+                .connect_options()
+                .get_filename()
+                .to_str()
+                .map(String::from),
+            port: None,
+            database: None,
+        };
+        Self { pool, attributes }
+    }
+}
+
+impl<DB: Database> PoolBuilder<DB> {
+    /// Set a custom name for the pool (for peer.service attribute).
+    pub fn with_name(mut self, name: impl Into<String>) -> Self {
+        self.attributes.name = Some(name.into());
+        self
+    }
+
+    /// Set the database name attribute.
+    pub fn with_database(mut self, database: impl Into<String>) -> Self {
+        self.attributes.database = Some(database.into());
+        self
+    }
+
+    /// Set the host attribute.
+    pub fn with_host(mut self, host: impl Into<String>) -> Self {
+        self.attributes.host = Some(host.into());
+        self
+    }
+
+    /// Set the port attribute.
+    pub fn with_port(mut self, port: u16) -> Self {
+        self.attributes.port = Some(port);
+        self
+    }
+
+    /// Build the [`Pool`] with the configured attributes.
+    pub fn build(self) -> Pool<DB> {
+        Pool {
+            inner: self.pool,
+            attributes: Arc::new(self.attributes),
+        }
+    }
+}
+
+/// An asynchronous pool of SQLx database connections with tracing instrumentation.
+///
+/// Wraps a SQLx [`Pool`] and propagates tracing attributes to all acquired connections.
+#[derive(Debug, Deref, DerefMut)]
+pub struct Pool<DB: Database> {
+    #[deref]
+    #[deref_mut]
+    inner: sqlx::Pool<DB>,
+    attributes: Arc<Attributes>,
+}
+
+// manually impl `Clone` because `DB` may not be `Clone`
+impl<DB: Database> Clone for Pool<DB> {
+    fn clone(&self) -> Self {
+        Self {
+            inner: self.inner.clone(),
+            attributes: self.attributes.clone(),
+        }
+    }
+}
+
+impl<DB> From<sqlx::Pool<DB>> for Pool<DB>
+where
+    DB: Database,
+    PoolBuilder<DB>: From<sqlx::Pool<DB>>,
+{
+    /// Convert a SQLx [`Pool`] into a tracing-instrumented [`Pool`].
+    fn from(inner: sqlx::Pool<DB>) -> Self {
+        PoolBuilder::from(inner).build()
+    }
+}
+
+/// Wrapper for a mutable SQLx connection reference with tracing attributes.
+///
+/// Used internally for transaction and pool connection executors.
+pub struct Connection<'c, DB: Database> {
+    inner: &'c mut DB::Connection,
+    attributes: Arc<Attributes>,
+}
+
+impl<'c, DB: Database> std::fmt::Debug for Connection<'c, DB> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Connection").finish_non_exhaustive()
+    }
+}
+
+/// A pooled SQLx connection instrumented for tracing.
+///
+/// Implements [`sqlx::Executor`] and propagates tracing attributes.
+#[derive(Debug)]
+pub struct PoolConnection<DB: Database> {
+    inner: sqlx::pool::PoolConnection<DB>,
+    attributes: Arc<Attributes>,
+}
+
+/// An in-progress database transaction or savepoint, instrumented for tracing.
+///
+/// Wraps a SQLx [`Transaction`] and propagates tracing attributes.
+#[derive(Debug)]
+pub struct Transaction<'c, DB: Database> {
+    inner: sqlx::Transaction<'c, DB>,
+    attributes: Arc<Attributes>,
+}
+
+/// Acquire connections or transactions from a database in a generic way.
+///
+/// Equivalent of [`sqlx::Acquire`] with tracing.
+pub trait Acquire<'c> {
+    type Database: Database;
+
+    fn acquire(
+        self,
+    ) -> BoxFuture<'c, Result<AnyConnection<'c, Self::Database>, sqlx::Error>>;
+
+    fn begin(
+        self,
+    ) -> BoxFuture<'c, Result<Transaction<'c, Self::Database>, sqlx::Error>>;
+}
+
+#[derive(Debug)]
+pub enum AnyConnection<'c, DB: Database> {
+    Pool(PoolConnection<DB>),
+    Raw(Connection<'c, DB>),
+}