add tracing support

Author: Protryon

tokio-postgres/Cargo.toml

@@ -58,6 +58,7 @@ postgres-protocol = { version = "0.6.5", path = "../postgres-protocol" }
 postgres-types = { version = "0.2.4", path = "../postgres-types" }
 tokio = { version = "1.27", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
+tracing = "0.1"
 rand = "0.8.5"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]

tokio-postgres/src/client.rs

@@ -9,12 +9,13 @@ use crate::simple_query::SimpleQueryStream;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
 use crate::tls::TlsConnect;
+use crate::trace::make_span;
 use crate::types::{Oid, ToSql, Type};
 #[cfg(feature = "runtime")]
 use crate::Socket;
 use crate::{
-    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, CopyInSink, Error,
-    Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
+    copy_in, copy_out, prepare, query, simple_query, slice_iter, CancelToken, Config, CopyInSink,
+    Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction, TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -34,6 +35,7 @@ use std::task::{Context, Poll};
 #[cfg(feature = "runtime")]
 use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite};
+use tracing::Instrument;
 
 pub struct Responses {
     receiver: mpsc::Receiver<BackendMessages>,
@@ -85,6 +87,11 @@ pub struct InnerClient {
     sender: mpsc::UnboundedSender<Request>,
     cached_typeinfo: Mutex<CachedTypeInfo>,
 
+    db_user: String,
+    db_name: String,
+    #[cfg(feature = "runtime")]
+    socket_config: Option<SocketConfig>,
+
     /// A buffer to use when writing out postgres commands.
     buffer: Mutex<BytesMut>,
 }
@@ -103,6 +110,18 @@ impl InnerClient {
         })
     }
 
+    pub fn db_user(&self) -> &str {
+        &self.db_user
+    }
+
+    pub fn db_name(&self) -> &str {
+        &self.db_name
+    }
+
+    pub(crate) fn socket_config(&self) -> Option<&SocketConfig> {
+        self.socket_config.as_ref()
+    }
+
     pub fn typeinfo(&self) -> Option<Statement> {
         self.cached_typeinfo.lock().typeinfo.clone()
     }
@@ -190,11 +209,16 @@ impl Client {
         ssl_mode: SslMode,
         process_id: i32,
         secret_key: i32,
+        config: &Config,
     ) -> Client {
         Client {
             inner: Arc::new(InnerClient {
                 sender,
                 cached_typeinfo: Default::default(),
+                db_user: config.user.clone().unwrap_or_default(),
+                db_name: config.dbname.clone().unwrap_or_default(),
+                #[cfg(feature = "runtime")]
+                socket_config: None,
                 buffer: Default::default(),
             }),
             #[cfg(feature = "runtime")]
@@ -211,6 +235,9 @@ impl Client {
 
     #[cfg(feature = "runtime")]
     pub(crate) fn set_socket_config(&mut self, socket_config: SocketConfig) {
+        if let Some(inner) = Arc::get_mut(&mut self.inner) {
+            inner.socket_config = Some(socket_config.clone());
+        }
         self.socket_config = Some(socket_config);
     }
 
@@ -231,7 +258,12 @@ impl Client {
         query: &str,
         parameter_types: &[Type],
     ) -> Result<Statement, Error> {
-        prepare::prepare(&self.inner, query, parameter_types).await
+        let span = make_span(&self.inner);
+        span.record("db.operation", "prepare");
+
+        prepare::prepare(&self.inner, query, parameter_types)
+            .instrument(span)
+            .await
     }
 
     /// Executes a statement, returning a vector of the resulting rows.

tokio-postgres/src/connect_raw.rs

@@ -101,7 +101,7 @@ where
     let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
 
     let (sender, receiver) = mpsc::unbounded();
-    let client = Client::new(sender, config.ssl_mode, process_id, secret_key);
+    let client = Client::new(sender, config.ssl_mode, process_id, secret_key, config);
     let connection = Connection::new(stream.inner, stream.delayed, parameters, receiver);
 
     Ok((client, connection))

tokio-postgres/src/lib.rs

@@ -176,6 +176,7 @@ mod socket;
 mod statement;
 pub mod tls;
 mod to_statement;
+mod trace;
 mod transaction;
 mod transaction_builder;
 pub mod types;

tokio-postgres/src/prepare.rs

@@ -2,6 +2,7 @@ use crate::client::InnerClient;
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
 use crate::error::SqlState;
+use crate::trace::make_span;
 use crate::types::{Field, Kind, Oid, Type};
 use crate::{query, slice_iter};
 use crate::{Column, Error, Statement};
@@ -15,6 +16,7 @@ use std::future::Future;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
+use tracing::Instrument;
 
 const TYPEINFO_QUERY: &str = "\
 SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, t.typbasetype, n.nspname, t.typrelid
@@ -100,15 +102,23 @@ pub async fn prepare(
         }
     }
 
-    Ok(Statement::new(client, name, parameters, columns))
+    Ok(Statement::new(
+        client,
+        name,
+        query.to_string(),
+        parameters,
+        columns,
+    ))
 }
 
 fn prepare_rec<'a>(
     client: &'a Arc<InnerClient>,
     query: &'a str,
     types: &'a [Type],
 ) -> Pin<Box<dyn Future<Output = Result<Statement, Error>> + 'a + Send>> {
-    Box::pin(prepare(client, query, types))
+    let span = make_span(client);
+    span.record("db.operation", "prepare");
+    Box::pin(prepare(client, query, types).instrument(span))
 }
 
 fn encode(client: &InnerClient, name: &str, query: &str, types: &[Type]) -> Result<Bytes, Error> {

tokio-postgres/src/query.rs

@@ -1,6 +1,7 @@
 use crate::client::{InnerClient, Responses};
 use crate::codec::FrontendMessage;
 use crate::connection::RequestMessages;
+use crate::trace::make_span;
 use crate::types::{BorrowToSql, IsNull};
 use crate::{Error, Portal, Row, Statement};
 use bytes::{Bytes, BytesMut};
@@ -13,6 +14,7 @@ use std::fmt;
 use std::marker::PhantomPinned;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use tracing::{Instrument, Span};
 
 struct BorrowToSqlParamsDebug<'a, T>(&'a [T]);
 
@@ -27,30 +29,74 @@ where
     }
 }
 
-pub async fn query<P, I>(
+fn encode_with_logs<P, I>(
     client: &InnerClient,
-    statement: Statement,
+    span: &Span,
+    statement: &Statement,
     params: I,
-) -> Result<RowStream, Error>
+) -> Result<Bytes, Error>
 where
     P: BorrowToSql,
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
 {
-    let buf = if log_enabled!(Level::Debug) {
+    if log_enabled!(Level::Debug) || !span.is_disabled() {
         let params = params.into_iter().collect::<Vec<_>>();
         debug!(
             "executing statement {} with parameters: {:?}",
             statement.name(),
             BorrowToSqlParamsDebug(params.as_slice()),
         );
-        encode(client, &statement, params)?
+        if !span.is_disabled() {
+            span.record("db.statement", statement.query());
+
+            // while OTEL supports arrays, we don't want to add that as a dependency, so debug view it is.
+            let raw_params = BorrowToSqlParamsDebug(params.as_slice());
+            span.record("db.statement.params", format!("{raw_params:?}"));
+        }
+        encode(client, statement, params)
     } else {
-        encode(client, &statement, params)?
-    };
-    let responses = start(client, buf).await?;
+        encode(client, statement, params)
+    }
+}
+
+async fn start_traced(client: &InnerClient, span: &Span, buf: Bytes) -> Result<Responses, Error> {
+    match start(client, buf).instrument(span.clone()).await {
+        Ok(response) => {
+            span.record("otel.status_code", "OK");
+            Ok(response)
+        }
+        Err(e) => {
+            span.record("otel.status_code", "ERROR");
+            span.record("exception.message", e.to_string());
+            Err(e)
+        }
+    }
+}
+
+pub async fn query<P, I>(
+    client: &InnerClient,
+    statement: Statement,
+    params: I,
+) -> Result<RowStream, Error>
+where
+    P: BorrowToSql,
+    I: IntoIterator<Item = P>,
+    I::IntoIter: ExactSizeIterator,
+{
+    let span = make_span(client);
+    span.record("db.operation", "query");
+
+    let buf = encode_with_logs(client, &span, &statement, params)?;
+
+    let responses = start_traced(client, &span, buf).await?;
+
+    span.in_scope(|| {
+        tracing::trace!("response ready");
+    });
     Ok(RowStream {
         statement,
+        span,
         responses,
         rows_affected: None,
         _p: PhantomPinned,
@@ -62,6 +108,10 @@ pub async fn query_portal(
     portal: &Portal,
     max_rows: i32,
 ) -> Result<RowStream, Error> {
+    let span = make_span(client);
+    span.record("db.statement", portal.statement().query());
+    span.record("db.operation", "portal");
+
     let buf = client.with_buf(|buf| {
         frontend::execute(portal.name(), max_rows, buf).map_err(Error::encode)?;
         frontend::sync(buf);
@@ -72,6 +122,7 @@ pub async fn query_portal(
 
     Ok(RowStream {
         statement: portal.statement().clone(),
+        span,
         responses,
         rows_affected: None,
         _p: PhantomPinned,
@@ -101,25 +152,25 @@ where
     I: IntoIterator<Item = P>,
     I::IntoIter: ExactSizeIterator,
 {
-    let buf = if log_enabled!(Level::Debug) {
-        let params = params.into_iter().collect::<Vec<_>>();
-        debug!(
-            "executing statement {} with parameters: {:?}",
-            statement.name(),
-            BorrowToSqlParamsDebug(params.as_slice()),
-        );
-        encode(client, &statement, params)?
-    } else {
-        encode(client, &statement, params)?
-    };
-    let mut responses = start(client, buf).await?;
+    let span = make_span(client);
+    span.record("db.operation", "execute");
+
+    let buf = encode_with_logs(client, &span, &statement, params)?;
+
+    let mut responses = start_traced(client, &span, buf).await?;
+
+    span.in_scope(|| {
+        tracing::trace!("response ready");
+    });
 
     let mut rows = 0;
     loop {
-        match responses.next().await? {
+        match responses.next().instrument(span.clone()).await? {
             Message::DataRow(_) => {}
             Message::CommandComplete(body) => {
                 rows = extract_row_affected(&body)?;
+                span.record("db.sql.rows_affected", rows);
+                tracing::trace!("execute complete");
             }
             Message::EmptyQueryResponse => rows = 0,
             Message::ReadyForQuery(_) => return Ok(rows),
@@ -206,6 +257,7 @@ pin_project! {
     /// A stream of table rows.
     pub struct RowStream {
         statement: Statement,
+        span: Span,
         responses: Responses,
         rows_affected: Option<u64>,
         #[pin]
@@ -218,13 +270,19 @@ impl Stream for RowStream {
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.project();
+        let _span = this.span.enter();
         loop {
             match ready!(this.responses.poll_next(cx)?) {
                 Message::DataRow(body) => {
                     return Poll::Ready(Some(Ok(Row::new(this.statement.clone(), body)?)))
                 }
                 Message::CommandComplete(body) => {
-                    *this.rows_affected = Some(extract_row_affected(&body)?);
+                    let rows_affected = extract_row_affected(&body)?;
+
+                    this.span.record("db.sql.rows_affected", rows_affected);
+                    tracing::trace!("query complete");
+
+                    *this.rows_affected = Some(rows_affected);
                 }
                 Message::EmptyQueryResponse | Message::PortalSuspended => {}
                 Message::ReadyForQuery(_) => return Poll::Ready(None),

tokio-postgres/src/statement.rs

@@ -11,6 +11,7 @@ use std::{
 struct StatementInner {
     client: Weak<InnerClient>,
     name: String,
+    query: String,
     params: Vec<Type>,
     columns: Vec<Column>,
 }
@@ -38,12 +39,14 @@ impl Statement {
     pub(crate) fn new(
         inner: &Arc<InnerClient>,
         name: String,
+        query: String,
         params: Vec<Type>,
         columns: Vec<Column>,
     ) -> Statement {
         Statement(Arc::new(StatementInner {
             client: Arc::downgrade(inner),
             name,
+            query,
             params,
             columns,
         }))
@@ -53,6 +56,11 @@ impl Statement {
         &self.0.name
     }
 
+    /// Returns the query that was used to create this statement.
+    pub fn query(&self) -> &str {
+        &self.0.query
+    }
+
     /// Returns the expected types of the statement's parameters.
     pub fn params(&self) -> &[Type] {
         &self.0.params