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
This commit is contained in:
aecsocket
2026-01-28 13:38:57 +00:00
committed by GitHub
parent 7cb7e881fa
commit e57c15b3ce
146 changed files with 7320 additions and 801 deletions

View File

@@ -0,0 +1,57 @@
on:
push:
branches: [main]
pull_request:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
cancel-in-progress: true
env:
RUSTFLAGS: '-Dwarnings'
name: check
jobs:
fmt:
runs-on: ubuntu-latest
name: fmt
steps:
- uses: actions/checkout@v5
- uses: dtolnay/rust-toolchain@stable
with:
components: rustfmt
- name: cargo fmt --check --all
run: cargo fmt --check --all
clippy:
runs-on: ubuntu-latest
name: clippy
steps:
- uses: actions/checkout@v5
- uses: dtolnay/rust-toolchain@stable
with:
components: clippy
- run: cargo clippy --all-features --all-targets --tests
test:
runs-on: ubuntu-latest
name: test
steps:
- uses: actions/checkout@v5
- uses: dtolnay/rust-toolchain@stable
- run: cargo test --workspace --all-features
features:
runs-on: ubuntu-latest
name: features
steps:
- uses: actions/checkout@v5
- uses: dtolnay/rust-toolchain@stable
- uses: taiki-e/install-action@v2
with:
tool: cargo-hack
- name: features powerset
run: cargo hack check --feature-powerset --tests

View File

@@ -0,0 +1,26 @@
name: release-plz
permissions:
pull-requests: write
contents: write
id-token: write
on:
push:
branches:
- main
jobs:
release-plz:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
- uses: dtolnay/rust-toolchain@stable
- name: Run release-plz
id: release-plz
uses: MarcoIeni/release-plz-action@v0.5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }}

1
packages/sqlx-tracing/.gitignore vendored Normal file
View File

@@ -0,0 +1 @@
/target

View File

@@ -0,0 +1,46 @@
# Changelog
All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
## [0.2.0](https://github.com/jdrouet/sqlx-tracing/compare/v0.1.0...v0.2.0) - 2025-10-02
### Added
- add attributes to pool
- make sure returned_rows is populated
- trace on pool connections and transactions
- make it work with PoolConnection
- make transaction part compile
- create pool-connection and transaction
### Fixed
- unused import
- create separate builder for sqlite and postgres
- please clippy
- remove unused traits
### Other
- use opentelemetry-testing from registry
- comment the code
- update readme with pool builder
- ensure pool queries are traced
- release v0.1.0
## [0.1.0](https://github.com/jdrouet/sqlx-tracing/releases/tag/v0.1.0) - 2025-09-07
### Other
- configure for auto release
- update cargo.toml
- set versions in dev deps
- add readme
- configure
- check that it works for sqlite and postgres
- simple project

3706
packages/sqlx-tracing/Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,35 @@
[package]
name = "sqlx-tracing"
version = "0.2.0"
edition = "2024"
description = "OpenTelemetry-compatible tracing for SQLx database operations in Rust."
documentation = "https://docs.rs/sqlx-tracing"
readme = "README.md"
homepage = "https://github.com/jdrouet/sqlx-tracing"
repository = "https://github.com/jdrouet/sqlx-tracing"
license = "MIT"
# authors = ["Jérémie Drouet <jeremie.drouet@gmail.com>"] # deprecated field, Tombi warns
keywords = ["database", "observability", "opentelemetry", "sqlx", "tracing"]
categories = [
"asynchronous",
"database",
"development-tools::debugging",
"development-tools::profiling"
]
[dependencies]
derive_more = { workspace = true, features = ["deref", "deref_mut"] }
futures = { version = "0.3" }
sqlx = { version = "0.8", default-features = false, features = ["derive"] }
tracing = { version = "0.1" }
[dev-dependencies]
opentelemetry = "0.30"
opentelemetry-testing = "0.1"
sqlx = { version = "0.8", features = ["runtime-tokio"] }
testcontainers = "0.25"
tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
[features]
postgres = ["sqlx/postgres"]
sqlite = ["sqlx/sqlite"]

View File

@@ -0,0 +1,91 @@
> [!NOTE]
>
> This is a vendored version of [`sqlx-tracing`](https://github.com/jdrouet/sqlx-tracing/) with modifications for our own purposes.
>
> This directory is licensed under the same license as the original project.
# sqlx-tracing
**sqlx-tracing** is a Rust library that provides OpenTelemetry-compatible tracing for SQLx database operations. It wraps SQLx connection pools and queries with tracing spans, enabling detailed observability of database interactions in distributed systems.
## Features
- **Automatic Tracing**: All SQLx queries executed through the provided pool are traced using [tracing](https://docs.rs/tracing) spans.
- **OpenTelemetry Integration**: Traces are compatible with OpenTelemetry, making it easy to export to collectors and observability platforms.
- **Error Recording**: Errors are automatically annotated with kind, message, and stacktrace in the tracing span.
- **Returned Rows**: The number of rows returned by queries is recorded for observability.
- **Database Agnostic**: Supports both PostgreSQL and SQLite via feature flags.
- **Macros**: Includes a macro for consistent span creation around queries.
## Usage
Add `sqlx-tracing` to your `Cargo.toml`:
```toml
[dependencies]
sqlx-tracing = "0.1"
sqlx = { version = "0.8", default-features = false, features = ["derive"] }
tracing = "0.1"
```
Enable the desired database feature:
- For PostgreSQL: `features = ["postgres"]`
- For SQLite: `features = ["sqlite"]`
Wrap your SQLx pool:
```rust,ignore
let pool = sqlx::PgPool::connect(&url).await?;
// the attributes will be resolved from the url
let traced_pool = sqlx_tracing::Pool::from(pool);
// or manually overwrite them
let traced_pool = sqlx_tracing::PoolBuilder::from(pool)
.with_name("my-domain-database")
.with_database("database")
.with_host("somewhere")
.with_port(1234)
.build();
```
Use the traced pool as you would a normal SQLx pool:
```rust,ignore
let result: Option<i32> = sqlx::query_scalar("select 1")
.fetch_optional(traced_pool)
.await?;
```
This works also with pool connections
```rust,ignore
let mut conn = traced_pool.acquire().await?;
let result: Option<i32> = sqlx::query_scalar("select 1")
.fetch_optional(&mut conn)
.await?;
```
And transactions
```rust,ignore
let mut tx = traced_pool.begin().await?;
let result: Option<i32> = sqlx::query_scalar("select 1")
.fetch_optional(&mut tx.executor())
.await?;
```
## OpenTelemetry Integration
To export traces, set up an OpenTelemetry collector and configure the tracing subscriber with the appropriate layers. See the `tests/common.rs` for a full example using `opentelemetry`, `opentelemetry-otlp`, and `tracing-opentelemetry`.
## Testing
Integration tests are provided for both PostgreSQL and SQLite, using [testcontainers](https://docs.rs/testcontainers) and a local OpenTelemetry collector.
## License
Licensed under MIT.
## Contributing
Contributions and issues are welcome! Please open a PR or issue on GitHub.

View File

@@ -0,0 +1,103 @@
use crate::{AnyConnection, Database};
impl<'c, 's, DB> sqlx::Executor<'s> for &'s mut AnyConnection<'c, DB>
where
DB: Database,
// I attempted to have `DB::ConnectionRef<'c>` unify to `&'c mut DB::Connection`.
// This *can* be unified apparently, but we can't actually use the fact that
// `DB::ConnectionRef<'c>: sqlx::Executor` if we do this.
// So, we need a casting function in `crate::Database`.
// Maybe this can be revisited sometime to not require the casting fn.
//
// for<'a> DB: Database<ConnectionRef<'a> = &'a mut <DB as sqlx::Database>::Connection>,
{
type Database = DB;
fn fetch_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<
sqlx::Either<
<Self::Database as sqlx::Database>::QueryResult,
<Self::Database as sqlx::Database>::Row,
>,
sqlx::Error,
>,
>
where
's: 'e,
E: 'q + sqlx::Execute<'q, Self::Database>,
{
match self {
AnyConnection::Pool(pool) => {
DB::cast_connection(&mut pool.inner).fetch_many(query)
}
AnyConnection::Raw(conn) => {
DB::cast_connection(conn.inner).fetch_many(query)
}
}
}
fn fetch_optional<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Option<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
's: 'e,
E: 'q + sqlx::Execute<'q, Self::Database>,
{
match self {
AnyConnection::Pool(pool) => {
DB::cast_connection(&mut pool.inner).fetch_optional(query)
}
AnyConnection::Raw(conn) => {
DB::cast_connection(conn.inner).fetch_optional(query)
}
}
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [<Self::Database as sqlx::Database>::TypeInfo],
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
's: 'e,
{
match self {
AnyConnection::Pool(pool) => DB::cast_connection(&mut pool.inner)
.prepare_with(sql, parameters),
AnyConnection::Raw(conn) => {
DB::cast_connection(conn.inner).prepare_with(sql, parameters)
}
}
}
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<sqlx::Describe<Self::Database>, sqlx::Error>,
>
where
's: 'e,
{
match self {
AnyConnection::Pool(pool) => {
DB::cast_connection(&mut pool.inner).describe(sql)
}
AnyConnection::Raw(conn) => {
DB::cast_connection(conn.inner).describe(sql)
}
}
}
}

View File

@@ -0,0 +1,512 @@
use futures::{StreamExt, TryStreamExt};
use tracing::Instrument;
impl<DB> AsMut<<DB as sqlx::Database>::Connection> for crate::PoolConnection<DB>
where
DB: crate::Database,
{
fn as_mut(&mut self) -> &mut <DB as sqlx::Database>::Connection {
self.inner.as_mut()
}
}
impl<'c, DB> sqlx::Executor<'c> for &'c mut crate::PoolConnection<DB>
where
DB: crate::Database,
// impl<'a> Executor<'a> for PgConnection
for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
{
type Database = DB;
#[doc(hidden)]
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<sqlx::Describe<Self::Database>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.describe", attrs, sql);
let fut = self.inner.as_mut().describe(sql);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute", attrs, sql);
let fut = self.inner.execute(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute_many", attrs, sql);
let stream = self.inner.execute_many(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch", attrs, sql);
let stream = self.inner.fetch(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_all<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Vec<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let fut = self.inner.fetch_all(query);
Box::pin(
async move {
fut.await
.inspect(|res| {
let span = tracing::Span::current();
span.record("db.response.returned_rows", res.len());
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<
sqlx::Either<
<Self::Database as sqlx::Database>::QueryResult,
<Self::Database as sqlx::Database>::Row,
>,
sqlx::Error,
>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let stream = self.inner.fetch_many(query);
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_one<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_one", attrs, sql);
let fut = self.inner.fetch_one(query);
Box::pin(
async move {
fut.await
.inspect(|_| {
tracing::Span::current()
.record("db.response.returned_rows", 1);
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_optional<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Option<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_optional", attrs, sql);
let fut = self.inner.fetch_optional(query);
Box::pin(
async move {
fut.await
.inspect(|res| {
tracing::Span::current().record(
"db.response.returned_rows",
if res.is_some() { 1 } else { 0 },
);
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn prepare<'e, 'q: 'e>(
self,
query: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare", attrs, query);
let fut = self.inner.prepare(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [<Self::Database as sqlx::Database>::TypeInfo],
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare_with", attrs, sql);
let fut = self.inner.prepare_with(sql, parameters);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
}
impl<'c, DB> sqlx::Executor<'c> for &'c mut crate::Connection<'c, DB>
where
DB: crate::Database,
for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
{
type Database = DB;
#[doc(hidden)]
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<sqlx::Describe<Self::Database>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.describe", attrs, sql);
let fut = self.inner.describe(sql);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute", attrs, sql);
let fut = self.inner.execute(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute_many", attrs, sql);
let stream = self.inner.execute_many(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch", attrs, sql);
let stream = self.inner.fetch(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_all<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Vec<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let fut = self.inner.fetch_all(query);
Box::pin(
async move {
fut.await
.inspect(|res| {
let span = tracing::Span::current();
span.record("db.response.returned_rows", res.len());
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<
sqlx::Either<
<Self::Database as sqlx::Database>::QueryResult,
<Self::Database as sqlx::Database>::Row,
>,
sqlx::Error,
>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let stream = self.inner.fetch_many(query);
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_one<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_one", attrs, sql);
let fut = self.inner.fetch_one(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_one)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_optional<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Option<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
'c: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_optional", attrs, sql);
let fut = self.inner.fetch_optional(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_optional)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn prepare<'e, 'q: 'e>(
self,
query: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare", attrs, query);
let fut = self.inner.prepare(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [<Self::Database as sqlx::Database>::TypeInfo],
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
'c: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare_with", attrs, sql);
let fut = self.inner.prepare_with(sql, parameters);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
}

View File

@@ -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>),
}

View File

@@ -0,0 +1,310 @@
use futures::{StreamExt, TryStreamExt, future::BoxFuture};
use tracing::Instrument;
use crate::AnyConnection;
impl<'c, DB> crate::Acquire<'c> for &'c crate::Pool<DB>
where
DB: crate::Database,
{
type Database = DB;
fn acquire(
self,
) -> BoxFuture<'c, Result<AnyConnection<'c, DB>, sqlx::Error>> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.acquire", attrs);
let fut = self.inner.acquire();
let fut = async move {
let conn = fut.await.inspect_err(crate::span::record_error)?;
let conn = crate::PoolConnection {
inner: conn,
attributes: self.attributes.clone(),
};
let conn = AnyConnection::Pool(conn);
Ok::<_, sqlx::Error>(conn)
};
Box::pin(fut.instrument(span))
}
fn begin(
self,
) -> BoxFuture<
'c,
Result<crate::Transaction<'c, Self::Database>, sqlx::Error>,
> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.begin", attrs);
let fut = self.inner.begin();
Box::pin(
async move {
let txn = fut.await.inspect_err(crate::span::record_error)?;
let txn = crate::Transaction {
inner: txn,
attributes: self.attributes.clone(),
};
Ok::<_, sqlx::Error>(txn)
}
.instrument(span),
)
}
}
impl<DB: crate::Database> crate::Pool<DB> {
/// Retrieves a connection and immediately begins a new transaction.
///
/// The returned [`Transaction`] is instrumented for tracing.
///
/// [`Transaction`]: crate::Transaction
pub async fn begin(
&self,
) -> Result<crate::Transaction<'static, DB>, sqlx::Error> {
self.inner.begin().await.map(|inner| crate::Transaction {
inner,
attributes: self.attributes.clone(),
})
}
/// Acquires a pooled connection, instrumented for tracing.
pub async fn acquire(
&self,
) -> Result<crate::PoolConnection<DB>, sqlx::Error> {
self.inner
.acquire()
.await
.map(|inner| crate::PoolConnection {
attributes: self.attributes.clone(),
inner,
})
}
}
impl<'p, DB> sqlx::Executor<'p> for &crate::Pool<DB>
where
DB: crate::Database,
for<'c> &'c mut DB::Connection: sqlx::Executor<'c, Database = DB>,
{
type Database = DB;
#[doc(hidden)]
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<sqlx::Describe<Self::Database>, sqlx::Error>,
> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.describe", attrs, sql);
let fut = self.inner.describe(sql);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute", attrs, sql);
let fut = self.inner.execute(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute_many", attrs, sql);
let stream = self.inner.execute_many(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch", attrs, sql);
let stream = self.inner.fetch(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_all<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Vec<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let fut = self.inner.fetch_all(query);
Box::pin(
async move {
fut.await
.inspect(|res| {
let span = tracing::Span::current();
span.record("db.response.returned_rows", res.len());
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<
sqlx::Either<
<Self::Database as sqlx::Database>::QueryResult,
<Self::Database as sqlx::Database>::Row,
>,
sqlx::Error,
>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let stream = self.inner.fetch_many(query);
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_one<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_one", attrs, sql);
let fut = self.inner.fetch_one(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_one)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_optional<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Option<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_optional", attrs, sql);
let fut = self.inner.fetch_optional(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_optional)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn prepare<'e, 'q: 'e>(
self,
query: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare", attrs, query);
let fut = self.inner.prepare(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [<Self::Database as sqlx::Database>::TypeInfo],
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare_with", attrs, sql);
let fut = self.inner.prepare_with(sql, parameters);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
}

View File

@@ -0,0 +1,23 @@
impl crate::Database for sqlx::Postgres {
const SYSTEM: &'static str = "postgresql";
type ConnectionRef<'a> = &'a mut sqlx::PgConnection;
fn cast_connection<'c>(
conn: &'c mut <Self as sqlx::Database>::Connection,
) -> Self::ConnectionRef<'c> {
conn
}
// fn cast_pool_connection<'c>(
// conn: &'c mut PoolConnection<Self>,
// ) -> Self::PoolConnection<'c> {
// &mut conn.inner
// }
// fn cast_raw_connection<'c>(
// conn: &'c mut <Self as sqlx::Database>::Connection,
// ) -> Self::RawConnection<'c> {
// conn
// }
}

View File

@@ -0,0 +1,88 @@
/// Macro to create a tracing span for a SQLx operation with OpenTelemetry-compatible fields.
///
/// - `$name`: The operation name (e.g., "sqlx.execute").
/// - `$statement`: The SQL statement being executed.
/// - `$attributes`: Connection or pool attributes for peer and db context.
///
/// This macro is used internally by the crate to instrument all major SQLx operations.
#[macro_export]
macro_rules! instrument {
($name:expr, $attributes:expr $(, $statement:expr)? ) => {
tracing::info_span!(
$name,
// Database name (if available)
"db.name" = $attributes.database,
// Operation type (filled by SQLx or left empty)
"db.operation" = ::tracing::field::Empty,
// The SQL query text
$( "db.query.text" = $statement, )?
// Number of affected rows (to be filled after execution)
"db.response.affected_rows" = ::tracing::field::Empty,
// Number of returned rows (to be filled after execution)
"db.response.returned_rows" = ::tracing::field::Empty,
// Status code of the response (to be filled after execution)
"db.response.status_code" = ::tracing::field::Empty,
// Table name (optional, left empty)
"db.sql.table" = ::tracing::field::Empty,
// Database system (e.g., "postgresql", "sqlite")
"db.system.name" = DB::SYSTEM,
// Error type, message, and stacktrace (to be filled on error)
"error.type" = ::tracing::field::Empty,
"error.message" = ::tracing::field::Empty,
"error.stacktrace" = ::tracing::field::Empty,
// Peer (server) host and port
"net.peer.name" = $attributes.host,
"net.peer.port" = $attributes.port,
// OpenTelemetry semantic fields
"otel.kind" = "client",
"otel.status_code" = ::tracing::field::Empty,
"otel.status_description" = ::tracing::field::Empty,
// Peer service name (if set)
"peer.service" = $attributes.name,
)
};
}
/// Records that a single row was returned in the current tracing span.
/// Used for fetch_one operations.
pub fn record_one<T>(_value: &T) {
let span = tracing::Span::current();
span.record("db.response.returned_rows", 1);
}
/// Records whether an optional row was returned in the current tracing span.
/// Used for fetch_optional operations.
pub fn record_optional<T>(value: &Option<T>) {
let span = tracing::Span::current();
span.record(
"db.response.returned_rows",
if value.is_some() { 1 } else { 0 },
);
}
/// Records error details in the current tracing span for a SQLx error.
/// Sets OpenTelemetry status and error fields for observability backends.
pub fn record_error(err: &sqlx::Error) {
let span = tracing::Span::current();
// Mark the span as an error for OpenTelemetry
span.record("otel.status_code", "error");
span.record("otel.status_description", err.to_string());
// Classify error type as client or server
match err {
sqlx::Error::ColumnIndexOutOfBounds { .. }
| sqlx::Error::ColumnDecode { .. }
| sqlx::Error::ColumnNotFound(_)
| sqlx::Error::Decode { .. }
| sqlx::Error::Encode { .. }
| sqlx::Error::RowNotFound
| sqlx::Error::TypeNotFound { .. } => {
span.record("error.type", "client");
}
_ => {
span.record("error.type", "server");
}
}
// Attach error message and stacktrace for debugging
span.record("error.message", err.to_string());
span.record("error.stacktrace", format!("{err:?}"));
}

View File

@@ -0,0 +1,11 @@
impl crate::Database for sqlx::Sqlite {
const SYSTEM: &'static str = "sqlite";
type ConnectionRef<'a> = &'a mut sqlx::SqliteConnection;
fn cast_connection<'c>(
conn: &'c mut <Self as sqlx::Database>::Connection,
) -> Self::ConnectionRef<'c> {
conn
}
}

View File

@@ -0,0 +1,334 @@
use futures::{StreamExt, TryStreamExt, future::BoxFuture};
use sqlx::{Acquire, Error};
use tracing::Instrument;
use crate::AnyConnection;
impl<'c, DB> crate::Transaction<'c, DB>
where
DB: crate::Database,
{
/// Returns a tracing-instrumented executor for this transaction.
///
/// This allows running queries with full span context and attributes.
pub fn executor(&mut self) -> crate::Connection<'_, DB> {
crate::Connection {
inner: &mut *self.inner,
attributes: self.attributes.clone(),
}
}
/// Commits this transaction or savepoint.
pub async fn commit(self) -> Result<(), Error> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.commit", attrs);
let fut = self.inner.commit();
fut.instrument(span).await
}
/// Aborts this transaction or savepoint.
pub async fn rollback(self) -> Result<(), Error> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.rollback", attrs);
let fut = self.inner.rollback();
fut.instrument(span).await
}
}
impl<'c, 't, DB> crate::Acquire<'t> for &'t mut crate::Transaction<'c, DB>
where
DB: crate::Database,
{
type Database = DB;
#[inline]
fn acquire(
self,
) -> BoxFuture<'t, Result<AnyConnection<'t, DB>, sqlx::Error>> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.acquire", attrs);
let fut = self.inner.acquire();
let fut = async move {
let conn = fut.await.inspect_err(crate::span::record_error)?;
let conn = crate::Connection {
inner: conn,
attributes: attrs.clone(),
};
let conn = AnyConnection::Raw(conn);
Ok(conn)
};
Box::pin(fut.instrument(span))
}
fn begin(
self,
) -> BoxFuture<
't,
Result<crate::Transaction<'t, Self::Database>, sqlx::Error>,
> {
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.begin", attrs);
let fut = self.inner.begin();
let fut = async move {
let txn = fut.await.inspect_err(crate::span::record_error)?;
let txn = crate::Transaction {
inner: txn,
attributes: attrs.clone(),
};
Ok(txn)
};
Box::pin(fut.instrument(span))
}
}
/// Implements `sqlx::Executor` for a mutable reference to a tracing-instrumented transaction.
///
/// Each method creates a tracing span for the SQL operation, attaches relevant attributes,
/// and records errors or row counts as appropriate for observability.
impl<'c, 't, DB> sqlx::Executor<'t> for &'t mut crate::Transaction<'c, DB>
where
DB: crate::Database,
for<'a> &'a mut DB::Connection: sqlx::Executor<'a, Database = DB>,
{
type Database = DB;
#[doc(hidden)]
fn describe<'e, 'q: 'e>(
self,
sql: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<sqlx::Describe<Self::Database>, sqlx::Error>,
>
where
't: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.describe", attrs, sql);
Box::pin(
async move {
let fut = (&mut self.inner).describe(sql);
fut.await.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn execute<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute", attrs, sql);
let fut = (&mut self.inner).execute(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn execute_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::QueryResult, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.execute_many", attrs, sql);
let stream = (&mut self.inner).execute_many(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch", attrs, sql);
let stream = (&mut self.inner).fetch(query);
use futures::StreamExt;
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_all<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Vec<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let fut = (&mut self.inner).fetch_all(query);
Box::pin(
async move {
fut.await
.inspect(|res| {
let span = tracing::Span::current();
span.record("db.response.returned_rows", res.len());
})
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_many<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::stream::BoxStream<
'e,
Result<
sqlx::Either<
<Self::Database as sqlx::Database>::QueryResult,
<Self::Database as sqlx::Database>::Row,
>,
sqlx::Error,
>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_all", attrs, sql);
let stream = (&mut self.inner).fetch_many(query);
Box::pin(
stream
.inspect(move |_| {
let _enter = span.enter();
})
.inspect_err(crate::span::record_error),
)
}
fn fetch_one<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Row, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_one", attrs, sql);
let fut = (&mut self.inner).fetch_one(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_one)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn fetch_optional<'e, 'q: 'e, E>(
self,
query: E,
) -> futures::future::BoxFuture<
'e,
Result<Option<<Self::Database as sqlx::Database>::Row>, sqlx::Error>,
>
where
E: 'q + sqlx::Execute<'q, Self::Database>,
't: 'e,
{
let sql = query.sql();
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.fetch_optional", attrs, sql);
let fut = (&mut self.inner).fetch_optional(query);
Box::pin(
async move {
fut.await
.inspect(crate::span::record_optional)
.inspect_err(crate::span::record_error)
}
.instrument(span),
)
}
fn prepare<'e, 'q: 'e>(
self,
query: &'q str,
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
't: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare", attrs, query);
let fut = (&mut self.inner).prepare(query);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
fn prepare_with<'e, 'q: 'e>(
self,
sql: &'q str,
parameters: &'e [<Self::Database as sqlx::Database>::TypeInfo],
) -> futures::future::BoxFuture<
'e,
Result<<Self::Database as sqlx::Database>::Statement<'q>, sqlx::Error>,
>
where
't: 'e,
{
let attrs = &self.attributes;
let span = crate::instrument!("sqlx.prepare_with", attrs, sql);
let fut = (&mut self.inner).prepare_with(sql, parameters);
Box::pin(
async move { fut.await.inspect_err(crate::span::record_error) }
.instrument(span),
)
}
}

View File

@@ -0,0 +1,87 @@
//! Test that valid uses of the API compile.
#![expect(dead_code, reason = "only here to check that the code compiles")]
use sqlx::Postgres;
async fn a(db: sqlx_tracing::Pool<Postgres>) {
let _conn: sqlx_tracing::PoolConnection<Postgres> =
db.acquire().await.unwrap();
}
async fn b<'a, E>(exec: E)
where
E: sqlx_tracing::Acquire<'a, Database = Postgres>,
{
let mut conn: sqlx_tracing::AnyConnection<Postgres> =
exec.acquire().await.unwrap();
// sqlx::query("SELECT 1").execute(&mut conn).await.unwrap();
sqlx::query("SELECT 1").execute(&mut conn).await.unwrap();
}
async fn c<'a, E>(exec: E)
where
E: sqlx_tracing::Executor<'a, Database = Postgres>,
{
sqlx::query("SELECT 1").execute(exec).await.unwrap();
}
pub async fn list_many<'a, E>(exec: E)
where
E: sqlx::Executor<'a, Database = Postgres>,
{
sqlx::query(
"
SELECT
id, enum_id, value, ordering,
metadata, created
FROM loader_field_enum_values
WHERE enum_id = ANY($1)
ORDER BY enum_id, ordering, created DESC
",
)
.fetch_all(exec)
.await
.unwrap();
}
async fn insert_sqlx(transaction: &mut sqlx::Transaction<'_, Postgres>) {
get_id_sqlx(&mut *transaction).await;
}
async fn insert<'t, 'c>(
transaction: &'t mut sqlx_tracing::Transaction<'c, Postgres>,
) {
get_id(&mut *transaction).await;
get_id(&mut *transaction).await;
sqlx::query("SELECT 1")
.execute(&mut *transaction)
.await
.unwrap();
sqlx::query("SELECT 1")
.execute(&mut *transaction)
.await
.unwrap();
}
async fn get_id_sqlx<'a, E>(_executor: E)
where
E: sqlx::Acquire<'a, Database = sqlx::Postgres>,
{
}
async fn get_id<'a, E>(_executor: E)
where
E: sqlx_tracing::Acquire<'a, Database = sqlx::Postgres>,
{
}
async fn d<'a, E>(exec: E)
where
E: sqlx_tracing::Acquire<'a, Database = Postgres>,
{
let mut exec = exec.acquire().await.unwrap();
sqlx::query("SELECT 1").fetch_one(&mut exec).await.unwrap();
sqlx::query("SELECT 1").fetch_one(&mut exec).await.unwrap();
}

View File

@@ -0,0 +1,47 @@
use std::time::Duration;
use opentelemetry::trace::{FutureExt, TraceContextExt, Tracer};
pub async fn should_trace<'c, DB, E>(
name: &'static str,
system: &'static str,
observability: &opentelemetry_testing::ObservabilityContainer,
provider: &opentelemetry_testing::OpenTelemetryProvider,
executor: E,
) where
DB: sqlx::Database,
E: sqlx::Executor<'c, Database = DB>,
for<'q> DB::Arguments<'q>: 'q + sqlx::IntoArguments<'q, DB>,
(i32,): Send + Unpin + for<'r> sqlx::FromRow<'r, DB::Row>,
{
let scope = format!("should_{name}_{system}");
let tracer = opentelemetry::global::tracer(scope.clone());
let span = tracer.span_builder(name).start(&tracer);
let ctx = opentelemetry::Context::new().with_span(span);
let result: Option<i32> = sqlx::query_scalar("select 1")
.fetch_optional(executor)
.with_context(ctx)
.await
.unwrap();
assert_eq!(result, Some(1));
provider.flush();
tokio::time::sleep(Duration::from_secs(1)).await;
let traces = observability.json_traces();
let scope_span = traces.find_scope_span(&scope).unwrap();
let entry = scope_span.first_span().unwrap();
assert_eq!(entry.name, name);
let next = traces
.find_child(&entry.span_id, "sqlx.fetch_optional")
.unwrap();
assert_eq!(next.string_attribute("db.system.name").unwrap(), system);
assert_eq!(next.string_attribute("db.query.text").unwrap(), "select 1");
assert_eq!(
next.int_attribute("db.response.returned_rows").unwrap(),
"1"
);
}

View File

@@ -0,0 +1,80 @@
#![cfg(feature = "postgres")]
use std::time::Duration;
use sqlx::Postgres;
use testcontainers::{
GenericImage, ImageExt,
core::{ContainerPort, WaitFor},
runners::AsyncRunner,
};
mod common;
#[derive(Debug)]
struct PostgresContainer {
container: testcontainers::ContainerAsync<testcontainers::GenericImage>,
}
impl PostgresContainer {
async fn create() -> Self {
let container = GenericImage::new("postgres", "15-alpine")
.with_wait_for(WaitFor::message_on_stderr(
"database system is ready to accept connections",
))
.with_exposed_port(ContainerPort::Tcp(5432))
.with_env_var("POSTGRES_USER", "postgres")
.with_env_var("POSTGRES_DB", "postgres")
.with_env_var("POSTGRES_HOST_AUTH_METHOD", "trust")
.with_startup_timeout(Duration::from_secs(60))
.start()
.await
.expect("starting a postgres database");
Self { container }
}
async fn client(&self) -> sqlx_tracing::Pool<Postgres> {
let port = self.container.get_host_port_ipv4(5432).await.unwrap();
let url = format!("postgres://postgres@localhost:{port}/postgres");
sqlx::PgPool::connect(&url)
.await
.map(sqlx_tracing::Pool::from)
.unwrap()
}
}
#[tokio::test]
async fn execute() {
let observability = opentelemetry_testing::ObservabilityContainer::create().await;
let provider = observability.install().await;
let container = PostgresContainer::create().await;
let pool = container.client().await;
common::should_trace("trace_pool", "postgresql", &observability, &provider, &pool).await;
{
let mut conn = pool.acquire().await.unwrap();
common::should_trace(
"trace_conn",
"postgresql",
&observability,
&provider,
&mut conn,
)
.await;
}
{
let mut tx: sqlx_tracing::Transaction<'_, Postgres> = pool.begin().await.unwrap();
common::should_trace(
"trace_tx",
"postgresql",
&observability,
&provider,
&mut tx.executor(),
)
.await;
}
}

View File

@@ -0,0 +1,33 @@
#![cfg(feature = "sqlite")]
use sqlx::Sqlite;
mod common;
#[tokio::test]
async fn execute() {
let observability = opentelemetry_testing::ObservabilityContainer::create().await;
let provider = observability.install().await;
let pool = sqlx::SqlitePool::connect(":memory:").await.unwrap();
let pool = sqlx_tracing::Pool::from(pool);
common::should_trace("trace_pool", "sqlite", &observability, &provider, &pool).await;
{
let mut conn = pool.acquire().await.unwrap();
common::should_trace("trace_conn", "sqlite", &observability, &provider, &mut conn).await;
}
{
let mut tx: sqlx_tracing::Transaction<'_, Sqlite> = pool.begin().await.unwrap();
common::should_trace(
"trace_tx",
"sqlite",
&observability,
&provider,
&mut tx.executor(),
)
.await;
}
}