From 58b6a29e2e75b5affbb33faa4a9ba227e48bf18e Mon Sep 17 00:00:00 2001
From: Markus Ast <m@rkusa.st>
Date: Sat, 2 Dec 2023 20:10:28 +0100
Subject: [PATCH] add tracing to query execution

---
 postgres/Cargo.toml    |  1 +
 postgres/src/future.rs | 39 +++++++++++++++++++++++++++------------
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/postgres/Cargo.toml b/postgres/Cargo.toml
index 651dcb9..35c8482 100644
--- a/postgres/Cargo.toml
+++ b/postgres/Cargo.toml
@@ -20,6 +20,7 @@ sqlm-postgres-macros = { path = "../postgres-macros", default-features = false }
 time = { version = "0.3", optional = true }
 tokio-postgres = "0.7"
 tokio-postgres-rustls = "0.10"
+tracing = "0.1"
 uuid = { version = "1.4", optional = true }
 
 [dev-dependencies]
diff --git a/postgres/src/future.rs b/postgres/src/future.rs
index 6f3f57a..79e93e6 100644
--- a/postgres/src/future.rs
+++ b/postgres/src/future.rs
@@ -2,6 +2,9 @@ use std::future::{Future, IntoFuture};
 use std::marker::PhantomData;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use std::time::Instant;
+
+use tracing::Instrument;
 
 use crate::query::Query;
 use crate::{Error, Sql};
@@ -15,22 +18,34 @@ where
     type IntoFuture = SqlFuture<'a, T>;
 
     fn into_future(self) -> Self::IntoFuture {
+        let span =
+            tracing::debug_span!("sql query", query = self.query, parameters = ?self.parameters);
+        let start = Instant::now();
+
         SqlFuture {
-            future: Box::pin(async move {
-                let mut i = 1;
-                loop {
-                    match T::query(&self).await {
-                        Ok(r) => return Ok(r),
-                        Err(Error::Postgres(err)) if err.is_closed() && i <= 5 => {
-                            // retry pool size + 1 times if connection is closed (might have
-                            // received a closed one from the connection pool)
-                            i += 1;
-                            continue;
+            future: Box::pin(
+                async move {
+                    let mut i = 1;
+                    loop {
+                        match T::query(&self).await {
+                            Ok(r) => {
+                                let elapsed = start.elapsed();
+                                tracing::debug!(?elapsed, "sql query finished");
+                                return Ok(r);
+                            }
+                            Err(Error::Postgres(err)) if err.is_closed() && i <= 5 => {
+                                // retry pool size + 1 times if connection is closed (might have
+                                // received a closed one from the connection pool)
+                                i += 1;
+                                tracing::debug!("retry due to connection closed error");
+                                continue;
+                            }
+                            Err(err) => return Err(err),
                         }
-                        Err(err) => return Err(err),
                     }
                 }
-            }),
+                .instrument(span),
+            ),
             marker: PhantomData,
         }
     }