From c73d475733d4a68ce714ded6d94c29983a13171c Mon Sep 17 00:00:00 2001 From: Hicham Azimani Date: Fri, 13 Nov 2026 12:19:49 +0100 Subject: [PATCH 1/2] feat: Implement `sqlx::Acquire` for `&Pool` and expose the wrapped pool Some SQLx APIs take nothing an instrumented `Pool` can provide today: `Migrator::run` is generic over `sqlx::Acquire`, and `PgListener::connect_with` accepts only a `sqlx::Pool`. - `&Pool` now implements `sqlx::Acquire`. Both methods acquire through `Pool::acquire`, so wait time, pending requests and timeouts are recorded. `acquire` returns the new `AcquiredConnection`, which dereferences to the backend connection and records `use_time` when dropped. `begin` returns a plain `sqlx::Transaction`, which cannot report its release, so that lease records no `use_time`. - `Pool::inner()` returns the wrapped `sqlx::Pool`. Work done through it is not instrumented. Queries run on the raw connection or pool are not instrumented; the docs, README and CHANGELOG say so. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 5 ++ Cargo.toml | 2 +- README.md | 16 ++++++ src/connection.rs | 64 +++++++++++++++++++++-- src/lib.rs | 2 +- src/pool.rs | 79 +++++++++++++++++++++++++++- tests/pool_metrics.rs | 118 ++++++++++++++++++++++++++++++++++++++++++ tests/postgres.rs | 25 +++++++++ 8 files changed, 304 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ce819d4..d4fdf71 100644 --- a/CHANGELOG.md +++ b/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`, 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. +- `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. + ### 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)). diff --git a/Cargo.toml b/Cargo.toml index 0da2108..7802a10 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/README.md b/README.md index 4d2182f..bdedb56 100644 --- a/README.md +++ b/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: diff --git a/src/connection.rs b/src/connection.rs index 7b53e1c..ecd148d 100644 --- a/src/connection.rs +++ b/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 { /// Moves with an owned connection into a transaction, preserving the complete lease duration. #[derive(Debug)] pub(crate) struct ConnectionUsage { - pub(crate) use_time: Arc>, + /// `None` once [`discard`](Self::discard)ed: the lease end is not observable. + pub(crate) use_time: Option>>, pub(crate) acquired_at: Instant, pub(crate) base_attrs: Vec, } +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 std::fmt::Debug for PoolConnection { impl PoolConnection { 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 { + inner: sqlx::pool::PoolConnection, + _usage: ConnectionUsage, +} + +impl AcquiredConnection { + pub(crate) fn new(connection: PoolConnection) -> Self { + let PoolConnection { inner, usage, .. } = connection; + Self { + inner, + _usage: usage, + } + } +} + +impl std::fmt::Debug for AcquiredConnection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AcquiredConnection").finish_non_exhaustive() + } +} + +impl Deref for AcquiredConnection { + type Target = DB::Connection; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl DerefMut for AcquiredConnection { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.inner + } +} diff --git a/src/lib.rs b/src/lib.rs index d9a0f53..ed46954 100644 --- a/src/lib.rs +++ b/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}; diff --git a/src/pool.rs b/src/pool.rs index a57d4c3..921d4ca 100644 --- a/src/pool.rs +++ b/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 Pool { 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 Pool { }) } + /// 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 { + &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 Pool { } } +/// 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 { + type Database = DB; + + type Connection = AcquiredConnection; + + fn acquire(self) -> BoxFuture<'a, Result> { + Box::pin(async move { Pool::acquire(self).await.map(AcquiredConnection::new) }) + } + + fn begin(self) -> BoxFuture<'a, Result, 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`]. /// diff --git a/tests/pool_metrics.rs b/tests/pool_metrics.rs index 663ef19..f071a7a 100644 --- a/tests/pool_metrics.rs +++ b/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 + ); +} diff --git a/tests/postgres.rs b/tests/postgres.rs index d848008..7d94a12 100644 --- a/tests/postgres.rs +++ b/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"); +} From b55a42bbfadb58eda68ab3e57c17815a6fe1b8d5 Mon Sep 17 00:00:00 2001 From: Hicham Azimani Date: Fri, 13 Nov 2026 12:20:11 +0100 Subject: [PATCH 2/2] docs: Link the CHANGELOG entries to #47 Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d4fdf71..6c5bfdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Added -- `impl sqlx::Acquire for &Pool`, 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. -- `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. +- `impl sqlx::Acquire for &Pool`, 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