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

sqlx-sqlite/src/connection/worker.rs

@@ -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 {