feat: Implement sqlx::Acquire for &Pool and expose the wrapped pool

CHANGELOG.md

@@ -6,6 +6,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ## [Unreleased]
 
+### Added
+
+- `impl sqlx::Acquire for &Pool<DB>`, so `SQLx` APIs generic over `Acquire` – `sqlx::migrate::Migrator::run` among them – accept an instrumented pool. Acquisition goes through `Pool::acquire` and records `db.client.connection.wait_time`, `pending_requests` and `timeouts`; `acquire` returns the new `AcquiredConnection`, which dereferences to the backend connection and records `use_time` when dropped. `begin` returns a plain `sqlx::Transaction` and records no `use_time`. Queries run through either are not instrumented ([#47](https://github.com/chmodas/sqlx-otel/pull/47)).
+- `Pool::inner()`, the wrapped `sqlx::Pool`, for APIs that accept nothing else, such as `sqlx::postgres::PgListener::connect_with`. Work done through it is not instrumented ([#47](https://github.com/chmodas/sqlx-otel/pull/47)).
+
 ### Fixed
 
 - Implicit pool queries and `Pool::begin()` now route through the instrumented `Pool::acquire()`, so `db.client.connection.wait_time`, `use_time`, `pending_requests`, and `timeouts` account for every acquisition rather than only explicit `acquire()` calls. A workload written as `query(..).fetch_one(&pool)` previously borrowed its connection straight from `sqlx::Pool` and could saturate the pool without any of its waits or timeouts reaching telemetry. Queries issued inside a transaction still acquire once, and connection use time now follows the lease through transaction commit, rollback, or drop ([#40](https://github.com/chmodas/sqlx-otel/pull/40)).

Cargo.toml

@@ -53,7 +53,7 @@ dtor = "1"
 opentelemetry_sdk = { version = "0.33", features = ["testing", "rt-tokio"] }
 proptest = "1"
 serial_test = "3"
-sqlx = { version = "0.9", features = ["macros", "sqlite", "postgres", "mysql", "mysql-rsa", "runtime-tokio"] }
+sqlx = { version = "0.9", features = ["macros", "migrate", "sqlite", "postgres", "mysql", "mysql-rsa", "runtime-tokio"] }
 testcontainers = "0.28"
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 

README.md

@@ -87,6 +87,22 @@ let pool = PoolBuilder::from(raw_pool)
 
 `network.protocol.name` defaults to the backend's wire protocol (`"postgresql"` for Postgres, `"mysql"` for MySQL, absent for SQLite); override only when the connection is tunnelled through a different application-layer protocol. `network.transport` is not inferred from the connect string – callers who want this attribute on spans / metrics must declare it explicitly so the value reflects the deployment configuration rather than a guess.
 
+## APIs that need a plain `SQLx` pool
+
+`&Pool` implements `sqlx::Acquire`, so APIs generic over it accept the instrumented pool directly:
+
+```rust
+migrator.run(&pool).await?;
+```
+
+The acquisition is measured (`wait_time`, `pending_requests`, `timeouts`, and `use_time` for `acquire`), but queries run on the acquired connection are not instrumented. `Acquire::begin` hands back a plain `sqlx::Transaction`, which records no `use_time`.
+
+For APIs that take nothing but a `sqlx::Pool`, `Pool::inner` returns the wrapped pool. Nothing done through it is instrumented:
+
+```rust
+let listener = sqlx::postgres::PgListener::connect_with(pool.inner()).await?;
+```
+
 ## Per-query annotations
 
 Because the library does not parse SQL, per-query attributes are the caller's responsibility:

src/connection.rs

@@ -1,3 +1,4 @@
+use std::ops::{Deref, DerefMut};
 use std::sync::Arc;
 use std::time::Instant;
 
@@ -40,15 +41,25 @@ pub struct PoolConnection<DB: sqlx::Database> {
 /// Moves with an owned connection into a transaction, preserving the complete lease duration.
 #[derive(Debug)]
 pub(crate) struct ConnectionUsage {
-    pub(crate) use_time: Arc<Histogram<f64>>,
+    /// `None` once [`discard`](Self::discard)ed: the lease end is not observable.
+    pub(crate) use_time: Option<Arc<Histogram<f64>>>,
     pub(crate) acquired_at: Instant,
     pub(crate) base_attrs: Vec<KeyValue>,
 }
 
+impl ConnectionUsage {
+    /// Give up on recording this lease, for a connection handed to a type that does not report
+    /// when it releases the connection.
+    pub(crate) fn discard(mut self) {
+        self.use_time = None;
+    }
+}
+
 impl Drop for ConnectionUsage {
     fn drop(&mut self) {
-        self.use_time
-            .record(self.acquired_at.elapsed().as_secs_f64(), &self.base_attrs);
+        if let Some(use_time) = &self.use_time {
+            use_time.record(self.acquired_at.elapsed().as_secs_f64(), &self.base_attrs);
+        }
     }
 }
 
@@ -61,3 +72,50 @@ impl<DB: sqlx::Database> std::fmt::Debug for PoolConnection<DB> {
 impl<DB: Database> PoolConnection<DB> {
     impl_with_annotations_mut!();
 }
+
+/// A pooled connection acquired through the [`sqlx::Acquire`] implementation on `&Pool`.
+///
+/// It exists so that `SQLx` APIs generic over [`sqlx::Acquire`] – such as
+/// `sqlx::migrate::Migrator::run` – accept an instrumented [`Pool`](crate::Pool). It dereferences
+/// to the backend connection, the type those APIs expect.
+///
+/// The acquisition is instrumented exactly like [`Pool::acquire`](crate::Pool::acquire)
+/// (`db.client.connection.wait_time`, `pending_requests`, `timeouts`), and
+/// `db.client.connection.use_time` is recorded when the value is dropped. Queries run on the
+/// dereferenced connection are **not** instrumented: no spans, no
+/// `db.client.operation.duration`. Use [`Pool::acquire`](crate::Pool::acquire) to run
+/// instrumented queries on a dedicated connection.
+pub struct AcquiredConnection<DB: sqlx::Database> {
+    inner: sqlx::pool::PoolConnection<DB>,
+    _usage: ConnectionUsage,
+}
+
+impl<DB: sqlx::Database> AcquiredConnection<DB> {
+    pub(crate) fn new(connection: PoolConnection<DB>) -> Self {
+        let PoolConnection { inner, usage, .. } = connection;
+        Self {
+            inner,
+            _usage: usage,
+        }
+    }
+}
+
+impl<DB: sqlx::Database> std::fmt::Debug for AcquiredConnection<DB> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AcquiredConnection").finish_non_exhaustive()
+    }
+}
+
+impl<DB: sqlx::Database> Deref for AcquiredConnection<DB> {
+    type Target = DB::Connection;
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<DB: sqlx::Database> DerefMut for AcquiredConnection<DB> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}

src/lib.rs

@@ -267,7 +267,7 @@ mod transaction;
 
 pub use annotations::{Annotated, AnnotatedMut, QueryAnnotations};
 pub use attributes::QueryTextMode;
-pub use connection::PoolConnection;
+pub use connection::{AcquiredConnection, PoolConnection};
 pub use database::Database;
 pub use pool::{Pool, PoolBuilder};
 pub use query_ext::{AnnotatedQuery, QueryAnnotateExt};

src/pool.rs

@@ -1,11 +1,13 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+use futures::future::BoxFuture;
+
 use opentelemetry_semantic_conventions::metric as semconv_metric;
 
 use crate::annotations::{Annotated, QueryAnnotations};
 use crate::attributes::{ConnectionAttributes, QueryTextMode};
-use crate::connection::PoolConnection;
+use crate::connection::{AcquiredConnection, PoolConnection};
 use crate::database::Database;
 use crate::metrics::Metrics;
 use crate::transaction::Transaction;
@@ -404,7 +406,7 @@ impl<DB: Database> Pool<DB> {
             inner,
             state: self.state.clone(),
             usage: crate::connection::ConnectionUsage {
-                use_time: self.use_time.clone(),
+                use_time: Some(self.use_time.clone()),
                 acquired_at: std::time::Instant::now(),
                 base_attrs: attrs,
             },
@@ -437,6 +439,31 @@ impl<DB: Database> Pool<DB> {
         })
     }
 
+    /// The wrapped `sqlx::Pool`, for `SQLx` APIs that accept nothing else, such as
+    /// `sqlx::postgres::PgListener::connect_with`.
+    ///
+    /// Work done through the returned pool is **not** instrumented: no spans, no
+    /// `db.client.operation.duration`, no `db.client.connection.wait_time` / `use_time` /
+    /// `pending_requests` / `timeouts`. Its connections are the ones this pool shares, so they
+    /// still count towards `db.client.connection.count`.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # #[cfg(feature = "sqlite")]
+    /// # async fn _doc() -> Result<(), sqlx::Error> {
+    /// # use sqlx_otel::PoolBuilder;
+    /// let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
+    /// let raw: &sqlx::SqlitePool = pool.inner();
+    /// assert!(!raw.is_closed());
+    /// # Ok(())
+    /// # }
+    /// ```
+    #[must_use]
+    pub fn inner(&self) -> &sqlx::Pool<DB> {
+        &self.inner
+    }
+
     /// Shut down the pool, waiting for all connections to be released.
     pub async fn close(&self) {
         self.inner.close().await;
@@ -503,6 +530,54 @@ impl<DB: Database> Pool<DB> {
     }
 }
 
+/// Lets `SQLx` APIs generic over [`sqlx::Acquire`] – `sqlx::migrate::Migrator::run` among them –
+/// take an instrumented pool.
+///
+/// Both methods acquire through [`Pool::acquire`], so `db.client.connection.wait_time`,
+/// `pending_requests` and `timeouts` are recorded. `acquire` returns an [`AcquiredConnection`],
+/// which records `db.client.connection.use_time` when dropped. `begin` returns a plain
+/// `sqlx::Transaction`, which does not report when it releases its connection, so that lease
+/// records no `use_time`; prefer the inherent [`Pool::begin`] where the concrete type is known.
+/// Queries run through either are not instrumented.
+///
+/// # Example
+///
+/// ```no_run
+/// # #[cfg(feature = "sqlite")]
+/// # async fn _doc() -> Result<(), sqlx::Error> {
+/// # use sqlx_otel::PoolBuilder;
+/// // A function generic over `Acquire`, as `Migrator::run` is.
+/// async fn ping<'a>(
+///     db: impl sqlx::Acquire<'a, Database = sqlx::Sqlite>,
+/// ) -> Result<(), sqlx::Error> {
+///     let mut conn = db.acquire().await?;
+///     sqlx::query("SELECT 1").execute(&mut *conn).await?;
+///     Ok(())
+/// }
+///
+/// let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await?).build();
+/// ping(&pool).await?;
+/// # Ok(())
+/// # }
+/// ```
+impl<'a, DB: Database> sqlx::Acquire<'a> for &'a Pool<DB> {
+    type Database = DB;
+
+    type Connection = AcquiredConnection<DB>;
+
+    fn acquire(self) -> BoxFuture<'a, Result<Self::Connection, sqlx::Error>> {
+        Box::pin(async move { Pool::acquire(self).await.map(AcquiredConnection::new) })
+    }
+
+    fn begin(self) -> BoxFuture<'a, Result<sqlx::Transaction<'a, DB>, sqlx::Error>> {
+        Box::pin(async move {
+            let PoolConnection { inner, usage, .. } = Pool::acquire(self).await?;
+            usage.discard();
+            sqlx::Transaction::begin(inner, None).await
+        })
+    }
+}
+
 /// Decrements `db.client.connection.pending_requests` on every exit path from
 /// [`Pool::acquire`].
 ///

tests/pool_metrics.rs

@@ -1083,3 +1083,121 @@ mod async_std_runtime {
         drop(conn);
     }
 }
+
+/// `&Pool` implements `sqlx::Acquire`: the acquisition is measured like `Pool::acquire`, the lease
+/// ends when the connection is dropped, and queries on the raw connection are not instrumented.
+#[tokio::test]
+#[serial]
+async fn acquire_trait_measures_the_lease() {
+    let tel = common::TestTelemetry::install();
+    let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await.unwrap()).build();
+    let waits_before = histogram_count(&tel, "db.client.connection.wait_time");
+    let uses_before = histogram_count(&tel, "db.client.connection.use_time");
+
+    let mut conn = sqlx::Acquire::acquire(&pool).await.unwrap();
+    sqlx::query("SELECT 1").execute(&mut *conn).await.unwrap();
+    // `histogram_count` drains the exporters, so the spans are read first.
+    assert!(
+        tel.spans().is_empty(),
+        "queries on the raw connection are not instrumented"
+    );
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.wait_time"),
+        waits_before + 1
+    );
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.use_time"),
+        uses_before,
+        "the lease is still open"
+    );
+    assert_eq!(pending_count(&tel), 0);
+
+    drop(conn);
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.use_time"),
+        uses_before + 1
+    );
+}
+
+/// `Acquire::begin` measures the acquisition but records no `use_time`: a plain
+/// `sqlx::Transaction` does not report when it releases the connection.
+#[tokio::test]
+#[serial]
+async fn acquire_trait_begin_measures_the_wait_only() {
+    let tel = common::TestTelemetry::install();
+    let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await.unwrap()).build();
+    let waits_before = histogram_count(&tel, "db.client.connection.wait_time");
+    let uses_before = histogram_count(&tel, "db.client.connection.use_time");
+
+    let mut tx = sqlx::Acquire::begin(&pool).await.unwrap();
+    sqlx::query("SELECT 1").execute(&mut *tx).await.unwrap();
+    tx.commit().await.unwrap();
+
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.wait_time"),
+        waits_before + 1
+    );
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.use_time"),
+        uses_before
+    );
+    assert_eq!(pending_count(&tel), 0);
+}
+
+/// The use case `Acquire` exists for: a `Migrator` runs on the instrumented pool.
+#[tokio::test]
+#[serial]
+async fn migrator_runs_on_the_instrumented_pool() {
+    use sqlx::migrate::{Migration, MigrationType, Migrator};
+    let tel = common::TestTelemetry::install();
+    let pool = PoolBuilder::from(sqlx::SqlitePool::connect(":memory:").await.unwrap()).build();
+    let waits_before = histogram_count(&tel, "db.client.connection.wait_time");
+
+    let migrator = Migrator::with_migrations(vec![Migration::new(
+        1,
+        "create t".into(),
+        MigrationType::Simple,
+        sqlx::SqlStr::from_static("CREATE TABLE t (id integer)"),
+        false,
+    )]);
+    migrator.run(&pool).await.unwrap();
+
+    let row: (i64,) = sqlx::query_as("SELECT count(*) FROM t")
+        .fetch_one(&pool)
+        .await
+        .unwrap();
+    assert_eq!(row.0, 0);
+    assert!(
+        histogram_count(&tel, "db.client.connection.wait_time") > waits_before,
+        "the migrator acquired through the instrumented pool"
+    );
+    assert_eq!(pending_count(&tel), 0);
+}
+
+/// `Pool::inner` is the wrapped pool: it shares the connections and is not instrumented.
+#[tokio::test]
+#[serial]
+async fn inner_pool_is_not_instrumented() {
+    let tel = common::TestTelemetry::install();
+    let raw = sqlx::sqlite::SqlitePoolOptions::new()
+        .max_connections(1)
+        .connect(":memory:")
+        .await
+        .unwrap();
+    let pool = PoolBuilder::from(raw).build();
+    pool.execute("CREATE TABLE t (id integer)").await.unwrap();
+    let waits_before = histogram_count(&tel, "db.client.connection.wait_time");
+
+    // One connection, one in-memory database: the table created through the wrapper is visible.
+    let row: (i64,) = sqlx::query_as("SELECT count(*) FROM t")
+        .fetch_one(pool.inner())
+        .await
+        .unwrap();
+    assert_eq!(row.0, 0);
+    // `histogram_count` drained the exporters: any span now came from the query above.
+    assert!(tel.spans().is_empty());
+    assert_eq!(
+        histogram_count(&tel, "db.client.connection.wait_time"),
+        waits_before
+    );
+}

tests/postgres.rs

@@ -1191,3 +1191,28 @@ async fn query_scalar_macro_fetch_all_with_annotations_via_pool() {
 async fn executor_side_query_side_parity_via_pool() {
     test_executor_side_query_side_parity_via_pool!(test_pool().await, common::POSTGRES_DIALECT);
 }
+
+// ===========================================================================
+// APIs that take a plain sqlx pool
+// ===========================================================================
+
+/// `PgListener::connect_with` accepts only a `sqlx::PgPool`: `Pool::inner` provides it.
+#[tokio::test]
+#[serial]
+async fn listener_connects_through_inner_pool() {
+    let pool = test_pool().await;
+    let mut listener = sqlx::postgres::PgListener::connect_with(pool.inner())
+        .await
+        .unwrap();
+    listener.listen("sqlx_otel_inner").await.unwrap();
+
+    pool.execute("NOTIFY sqlx_otel_inner, 'ping'")
+        .await
+        .unwrap();
+
+    let notification = tokio::time::timeout(Duration::from_secs(5), listener.recv())
+        .await
+        .expect("no notification")
+        .unwrap();
+    assert_eq!(notification.payload(), "ping");
+}