fix: spans in sqlite tracing (#2876) (#3056)

This commit is contained in:
zoomiti 2024-02-15 22:27:22 -08:00 committed by GitHub
parent 4c057a0628
commit a1e4984c6c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -13,6 +13,7 @@ use sqlx_core::transaction::{
begin_ansi_transaction_sql, commit_ansi_transaction_sql, rollback_ansi_transaction_sql,
};
use sqlx_core::Either;
use tracing::span::Span;
use crate::connection::describe::describe;
use crate::connection::establish::EstablishParams;
@ -27,7 +28,7 @@ use crate::{Sqlite, SqliteArguments, SqliteQueryResult, SqliteRow, SqliteStateme
// unlikely.
pub(crate) struct ConnectionWorker {
command_tx: flume::Sender<Command>,
command_tx: flume::Sender<(Command, tracing::Span)>,
/// The `sqlite3` pointer. NOTE: access is unsynchronized!
pub(crate) _handle_raw: ConnectionHandleRaw,
/// Mutex for locking access to the database.
@ -117,7 +118,8 @@ impl ConnectionWorker {
// would rollback an already completed transaction.
let mut ignore_next_start_rollback = false;
for cmd in command_rx {
for (cmd, span) in command_rx {
let _guard = span.enter();
match cmd {
Command::Prepare { query, tx } => {
tx.send(prepare(&mut conn, &query).map(|prepared| {
@ -289,12 +291,15 @@ impl ConnectionWorker {
let (tx, rx) = flume::bounded(chan_size);
self.command_tx
.send_async(Command::Execute {
query: query.into(),
arguments: args.map(SqliteArguments::into_static),
persistent,
tx,
})
.send_async((
Command::Execute {
query: query.into(),
arguments: args.map(SqliteArguments::into_static),
persistent,
tx,
},
Span::current(),
))
.await
.map_err(|_| Error::WorkerCrashed)?;
@ -318,7 +323,7 @@ impl ConnectionWorker {
pub(crate) fn start_rollback(&mut self) -> Result<(), Error> {
self.command_tx
.send(Command::Rollback { tx: None })
.send((Command::Rollback { tx: None }, Span::current()))
.map_err(|_| Error::WorkerCrashed)
}
@ -333,7 +338,7 @@ impl ConnectionWorker {
let (tx, rx) = oneshot::channel();
self.command_tx
.send_async(command(tx))
.send_async((command(tx), Span::current()))
.await
.map_err(|_| Error::WorkerCrashed)?;
@ -347,7 +352,7 @@ impl ConnectionWorker {
let (tx, rx) = rendezvous_oneshot::channel();
self.command_tx
.send_async(command(tx))
.send_async((command(tx), Span::current()))
.await
.map_err(|_| Error::WorkerCrashed)?;
@ -362,7 +367,8 @@ impl ConnectionWorker {
let (guard, res) = futures_util::future::join(
// we need to join the wait queue for the lock before we send the message
self.shared.conn.lock(),
self.command_tx.send_async(Command::UnlockDb),
self.command_tx
.send_async((Command::UnlockDb, Span::current())),
)
.await;
@ -379,7 +385,7 @@ impl ConnectionWorker {
let send_res = self
.command_tx
.send(Command::Shutdown { tx })
.send((Command::Shutdown { tx }, Span::current()))
.map_err(|_| Error::WorkerCrashed);
async move {