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:
57
packages/sqlx-tracing/.github/workflows/main.yml
vendored
Normal file
57
packages/sqlx-tracing/.github/workflows/main.yml
vendored
Normal 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
|
||||
26
packages/sqlx-tracing/.github/workflows/release-pr.yml
vendored
Normal file
26
packages/sqlx-tracing/.github/workflows/release-pr.yml
vendored
Normal 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
1
packages/sqlx-tracing/.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
/target
|
||||
46
packages/sqlx-tracing/CHANGELOG.md
Normal file
46
packages/sqlx-tracing/CHANGELOG.md
Normal 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
3706
packages/sqlx-tracing/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load Diff
35
packages/sqlx-tracing/Cargo.toml
Normal file
35
packages/sqlx-tracing/Cargo.toml
Normal 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"]
|
||||
91
packages/sqlx-tracing/README.md
Normal file
91
packages/sqlx-tracing/README.md
Normal 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.
|
||||
103
packages/sqlx-tracing/src/any_connection.rs
Normal file
103
packages/sqlx-tracing/src/any_connection.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
512
packages/sqlx-tracing/src/connection.rs
Normal file
512
packages/sqlx-tracing/src/connection.rs
Normal 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),
|
||||
)
|
||||
}
|
||||
}
|
||||
219
packages/sqlx-tracing/src/lib.rs
Normal file
219
packages/sqlx-tracing/src/lib.rs
Normal 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>),
|
||||
}
|
||||
310
packages/sqlx-tracing/src/pool.rs
Normal file
310
packages/sqlx-tracing/src/pool.rs
Normal 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),
|
||||
)
|
||||
}
|
||||
}
|
||||
23
packages/sqlx-tracing/src/postgres.rs
Normal file
23
packages/sqlx-tracing/src/postgres.rs
Normal 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
|
||||
// }
|
||||
}
|
||||
88
packages/sqlx-tracing/src/span.rs
Normal file
88
packages/sqlx-tracing/src/span.rs
Normal 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:?}"));
|
||||
}
|
||||
11
packages/sqlx-tracing/src/sqlite.rs
Normal file
11
packages/sqlx-tracing/src/sqlite.rs
Normal 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
|
||||
}
|
||||
}
|
||||
334
packages/sqlx-tracing/src/transaction.rs
Normal file
334
packages/sqlx-tracing/src/transaction.rs
Normal 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),
|
||||
)
|
||||
}
|
||||
}
|
||||
87
packages/sqlx-tracing/tests/api.rs
Normal file
87
packages/sqlx-tracing/tests/api.rs
Normal 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();
|
||||
}
|
||||
47
packages/sqlx-tracing/tests/common.rs
Normal file
47
packages/sqlx-tracing/tests/common.rs
Normal 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"
|
||||
);
|
||||
}
|
||||
80
packages/sqlx-tracing/tests/postgres.rs
Normal file
80
packages/sqlx-tracing/tests/postgres.rs
Normal 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;
|
||||
}
|
||||
}
|
||||
33
packages/sqlx-tracing/tests/sqlite.rs
Normal file
33
packages/sqlx-tracing/tests/sqlite.rs
Normal 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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user