sqlx_tracing/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use std::sync::Arc;
4
5mod connection;
6mod pool;
7pub mod prelude;
8pub(crate) mod span;
9mod transaction;
10
11#[cfg(feature = "postgres")]
12pub mod postgres;
13
14#[cfg(feature = "sqlite")]
15pub mod sqlite;
16
17/// Attributes describing the database connection and context.
18/// Used for span enrichment and attribute propagation.
19#[derive(Debug, Default)]
20struct Attributes {
21    name: Option<String>,
22    host: Option<String>,
23    port: Option<u16>,
24    database: Option<String>,
25}
26
27/// Builder for constructing a [`Pool`] with custom attributes.
28///
29/// Allows setting database name, host, port, and other identifying information
30/// for tracing purposes.
31#[derive(Debug)]
32pub struct PoolBuilder<DB: sqlx::Database> {
33    pool: sqlx::Pool<DB>,
34    attributes: Attributes,
35}
36
37// this is required because `pool.connect_options().to_url_lossy()` panics with sqlite
38#[cfg(feature = "postgres")]
39impl From<sqlx::Pool<sqlx::Postgres>> for PoolBuilder<sqlx::Postgres> {
40    /// Create a new builder from an existing SQLx pool.
41    fn from(pool: sqlx::Pool<sqlx::Postgres>) -> Self {
42        use sqlx::ConnectOptions;
43
44        let url = pool.connect_options().to_url_lossy();
45        let attributes = Attributes {
46            name: None,
47            host: url.host_str().map(String::from),
48            port: url.port(),
49            database: url
50                .path_segments()
51                .and_then(|mut segments| segments.next().map(String::from)),
52        };
53        Self { pool, attributes }
54    }
55}
56
57// this is required because `pool.connect_options().to_url_lossy()` panics with sqlite
58#[cfg(feature = "sqlite")]
59impl From<sqlx::Pool<sqlx::Sqlite>> for PoolBuilder<sqlx::Sqlite> {
60    /// Create a new builder from an existing SQLx pool.
61    fn from(pool: sqlx::Pool<sqlx::Sqlite>) -> Self {
62        let attributes = Attributes {
63            name: None,
64            host: pool
65                .connect_options()
66                .get_filename()
67                .to_str()
68                .map(String::from),
69            port: None,
70            database: None,
71        };
72        Self { pool, attributes }
73    }
74}
75
76impl<DB: sqlx::Database> PoolBuilder<DB> {
77    /// Set a custom name for the pool (for peer.service attribute).
78    pub fn with_name(mut self, name: impl Into<String>) -> Self {
79        self.attributes.name = Some(name.into());
80        self
81    }
82
83    /// Set the database name attribute.
84    pub fn with_database(mut self, database: impl Into<String>) -> Self {
85        self.attributes.database = Some(database.into());
86        self
87    }
88
89    /// Set the host attribute.
90    pub fn with_host(mut self, host: impl Into<String>) -> Self {
91        self.attributes.host = Some(host.into());
92        self
93    }
94
95    /// Set the port attribute.
96    pub fn with_port(mut self, port: u16) -> Self {
97        self.attributes.port = Some(port);
98        self
99    }
100
101    /// Build the [`Pool`] with the configured attributes.
102    pub fn build(self) -> Pool<DB> {
103        Pool {
104            inner: self.pool,
105            attributes: Arc::new(self.attributes),
106        }
107    }
108}
109
110/// An asynchronous pool of SQLx database connections with tracing instrumentation.
111///
112/// Wraps a SQLx [`Pool`] and propagates tracing attributes to all acquired connections.
113#[derive(Clone, Debug)]
114pub struct Pool<DB>
115where
116    DB: sqlx::Database,
117{
118    inner: sqlx::Pool<DB>,
119    attributes: Arc<Attributes>,
120}
121
122impl<DB> From<sqlx::Pool<DB>> for Pool<DB>
123where
124    DB: sqlx::Database,
125    PoolBuilder<DB>: From<sqlx::Pool<DB>>,
126{
127    /// Convert a SQLx [`Pool`] into a tracing-instrumented [`Pool`].
128    fn from(inner: sqlx::Pool<DB>) -> Self {
129        PoolBuilder::from(inner).build()
130    }
131}
132
133impl<DB> Pool<DB>
134where
135    DB: sqlx::Database,
136{
137    /// Retrieves a connection and immediately begins a new transaction.
138    ///
139    /// The returned [`Transaction`] is instrumented for tracing.
140    pub async fn begin<'c>(&'c self) -> Result<Transaction<'c, DB>, sqlx::Error> {
141        self.inner.begin().await.map(|inner| Transaction {
142            inner,
143            attributes: self.attributes.clone(),
144        })
145    }
146
147    /// Acquires a pooled connection, instrumented for tracing.
148    pub async fn acquire(&self) -> Result<PoolConnection<DB>, sqlx::Error> {
149        self.inner.acquire().await.map(|inner| PoolConnection {
150            attributes: self.attributes.clone(),
151            inner,
152        })
153    }
154}
155
156/// Wrapper for a mutable SQLx connection reference with tracing attributes.
157///
158/// Used internally for transaction and pool connection executors.
159pub struct Connection<'c, DB>
160where
161    DB: sqlx::Database,
162{
163    inner: &'c mut DB::Connection,
164    attributes: Arc<Attributes>,
165}
166
167impl<'c, DB: sqlx::Database> std::fmt::Debug for Connection<'c, DB> {
168    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
169        f.debug_struct("Connection").finish_non_exhaustive()
170    }
171}
172
173/// A pooled SQLx connection instrumented for tracing.
174///
175/// Implements [`sqlx::Executor`] and propagates tracing attributes.
176#[derive(Debug)]
177pub struct PoolConnection<DB>
178where
179    DB: sqlx::Database,
180{
181    inner: sqlx::pool::PoolConnection<DB>,
182    attributes: Arc<Attributes>,
183}
184
185/// An in-progress database transaction or savepoint, instrumented for tracing.
186///
187/// Wraps a SQLx [`Transaction`] and propagates tracing attributes.
188#[derive(Debug)]
189pub struct Transaction<'c, DB>
190where
191    DB: sqlx::Database,
192{
193    inner: sqlx::Transaction<'c, DB>,
194    attributes: Arc<Attributes>,
195}
sqlx_tracing/lib.rs

sqlx_tracing/
lib.rs
