Fix transction rollback on Future drop

postgres-protocol/src/authentication/sasl.rs

@@ -4,7 +4,7 @@ use base64::display::Base64Display;
 use base64::engine::general_purpose::STANDARD;
 use base64::Engine;
 use hmac::{Hmac, Mac};
-use rand::{self, Rng};
+use rand::Rng;
 use sha2::digest::FixedOutput;
 use sha2::{Digest, Sha256};
 use std::fmt::Write;

postgres-types/Cargo.toml

@@ -13,6 +13,7 @@ categories = ["database"]
 [features]
 derive = ["postgres-derive"]
 array-impls = ["array-init"]
+tuple-impls = []
 with-bit-vec-0_6 = ["bit-vec-06"]
 with-cidr-0_2 = ["cidr-02"]
 with-chrono-0_4 = ["chrono-04"]

postgres-types/src/chrono_04.rs

@@ -111,9 +111,11 @@ impl ToSql for DateTime<FixedOffset> {
 impl<'a> FromSql<'a> for NaiveDate {
     fn from_sql(_: &Type, raw: &[u8]) -> Result<NaiveDate, Box<dyn Error + Sync + Send>> {
         let jd = types::date_from_sql(raw)?;
+        let jd = Duration::try_days(i64::from(jd))
+            .ok_or_else(|| "value too large to decode")?;
         base()
             .date()
-            .checked_add_signed(Duration::days(i64::from(jd)))
+            .checked_add_signed(jd)
             .ok_or_else(|| "value too large to decode".into())
     }
 

postgres-types/src/lib.rs

@@ -642,6 +642,62 @@ impl<'a, T: FromSql<'a>, const N: usize> FromSql<'a> for [T; N] {
     }
 }
 
+macro_rules! impl_from_sql_tuple {
+    ($n:expr; $($ty_ident:ident),*; $($var_ident:ident),*) => {
+        impl<'a, $($ty_ident),*> FromSql<'a> for ($($ty_ident,)*)
+        where
+            $($ty_ident: FromSql<'a>),*
+        {
+            fn from_sql(
+                _: &Type,
+                mut raw: &'a [u8],
+            ) -> Result<($($ty_ident,)*), Box<dyn Error + Sync + Send>> {
+                let num_fields = private::read_be_i32(&mut raw)?;
+                if num_fields as usize != $n {
+                    return Err(format!(
+                        "Postgres record field count does not match Rust tuple length: {} vs {}",
+                        num_fields,
+                        $n,
+                    ).into());
+                }
+
+                $(
+                    let oid = private::read_be_i32(&mut raw)? as u32;
+                    let ty = match Type::from_oid(oid) {
+                        None => {
+                            return Err(format!(
+                                "cannot decode OID {} inside of anonymous record",
+                                oid,
+                            ).into());
+                        }
+                        Some(ty) if !$ty_ident::accepts(&ty) => {
+                            return Err(Box::new(WrongType::new::<$ty_ident>(ty.clone())));
+                        }
+                        Some(ty) => ty,
+                    };
+                    let $var_ident = private::read_value(&ty, &mut raw)?;
+                )*
+
+                Ok(($($var_ident,)*))
+            }
+
+            fn accepts(ty: &Type) -> bool {
+                match ty.kind() {
+                    Kind::Pseudo => *ty == Type::RECORD,
+                    Kind::Composite(fields) => fields.len() == $n,
+                    _ => false,
+                }
+            }
+        }
+    };
+}
+
+impl_from_sql_tuple!(0; ; );
+impl_from_sql_tuple!(1; T0; v0);
+impl_from_sql_tuple!(2; T0, T1; v0, v1);
+impl_from_sql_tuple!(3; T0, T1, T2; v0, v1, v2);
+impl_from_sql_tuple!(4; T0, T1, T2, T3; v0, v1, v2, v3);
+
 impl<'a, T: FromSql<'a>> FromSql<'a> for Box<[T]> {
     fn from_sql(ty: &Type, raw: &'a [u8]) -> Result<Self, Box<dyn Error + Sync + Send>> {
         Vec::<T>::from_sql(ty, raw).map(Vec::into_boxed_slice)

tokio-postgres/Cargo.toml

@@ -28,6 +28,7 @@ default = ["runtime"]
 runtime = ["tokio/net", "tokio/time"]
 
 array-impls = ["postgres-types/array-impls"]
+tuple-impls = ["postgres-types/tuple-impls"]
 with-bit-vec-0_6 = ["postgres-types/with-bit-vec-0_6"]
 with-chrono-0_4 = ["postgres-types/with-chrono-0_4"]
 with-eui48-0_4 = ["postgres-types/with-eui48-0_4"]

tokio-postgres/src/client.rs

@@ -13,8 +13,9 @@ 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, transaction::Savepoint,
+    CancelToken, CopyInSink, Error, Row, SimpleQueryMessage, Statement, ToStatement, Transaction,
+    TransactionBuilder,
 };
 use bytes::{Buf, BytesMut};
 use fallible_iterator::FallibleIterator;
@@ -469,8 +470,17 @@ impl Client {
     ///
     /// The transaction will roll back by default - use the `commit` method to commit it.
     pub async fn transaction(&mut self) -> Result<Transaction<'_>, Error> {
+        self.build_transaction().start().await
+    }
+
+    pub(crate) async fn start_transaction_with_rollback(
+        &mut self,
+        query: &str,
+        savepoint: Option<Savepoint>,
+    ) -> Result<Transaction<'_>, Error> {
         struct RollbackIfNotDone<'me> {
             client: &'me Client,
+            savepoint: Option<&'me Savepoint>,
             done: bool,
         }
 
@@ -480,8 +490,13 @@ impl Client {
                     return;
                 }
 
+                let query = if let Some(sp) = self.savepoint {
+                    format!("ROLLBACK TO {}", sp.name)
+                } else {
+                    "ROLLBACK".to_string()
+                };
                 let buf = self.client.inner().with_buf(|buf| {
-                    frontend::query("ROLLBACK", buf).unwrap();
+                    frontend::query(&query, buf).unwrap();
                     buf.split().freeze()
                 });
                 let _ = self
@@ -499,13 +514,14 @@ impl Client {
         {
             let mut cleaner = RollbackIfNotDone {
                 client: self,
+                savepoint: savepoint.as_ref(),
                 done: false,
             };
-            self.batch_execute("BEGIN").await?;
+            self.batch_execute(query).await?;
             cleaner.done = true;
         }
 
-        Ok(Transaction::new(self))
+        Ok(Transaction::new(self, savepoint))
     }
 
     /// Returns a builder for a transaction with custom settings.

tokio-postgres/src/transaction.rs

@@ -28,8 +28,8 @@ pub struct Transaction<'a> {
 }
 
 /// A representation of a PostgreSQL database savepoint.
-struct Savepoint {
-    name: String,
+pub(crate) struct Savepoint {
+    pub(crate) name: String,
     depth: u32,
 }
 
@@ -56,10 +56,10 @@ impl<'a> Drop for Transaction<'a> {
 }
 
 impl<'a> Transaction<'a> {
-    pub(crate) fn new(client: &'a mut Client) -> Transaction<'a> {
+    pub(crate) fn new(client: &'a mut Client, savepoint: Option<Savepoint>) -> Transaction<'a> {
         Transaction {
             client,
-            savepoint: None,
+            savepoint,
             done: false,
         }
     }
@@ -298,13 +298,11 @@ impl<'a> Transaction<'a> {
         let depth = self.savepoint.as_ref().map_or(0, |sp| sp.depth) + 1;
         let name = name.unwrap_or_else(|| format!("sp_{}", depth));
         let query = format!("SAVEPOINT {}", name);
-        self.batch_execute(&query).await?;
+        let savepoint = Savepoint { name, depth };
 
-        Ok(Transaction {
-            client: self.client,
-            savepoint: Some(Savepoint { name, depth }),
-            done: false,
-        })
+        self.client
+            .start_transaction_with_rollback(&query, Some(savepoint))
+            .await
     }
 
     /// Returns a reference to the underlying `Client`.

tokio-postgres/src/transaction_builder.rs

@@ -106,8 +106,8 @@ impl<'a> TransactionBuilder<'a> {
             query.push_str(s);
         }
 
-        self.client.batch_execute(&query).await?;
-
-        Ok(Transaction::new(self.client))
+        self.client
+            .start_transaction_with_rollback(&query, None)
+            .await
     }
 }

tokio-postgres/tests/test/types/mod.rs

@@ -379,6 +379,58 @@ async fn test_array_array_params() {
     .await;
 }
 
+#[tokio::test]
+async fn tuples() {
+    let client = connect("user=postgres").await;
+
+    let row = client.query_one("SELECT ROW()", &[]).await.unwrap();
+    row.get::<_, ()>(0);
+
+    let row = client.query_one("SELECT ROW(1)", &[]).await.unwrap();
+    let val: (i32,) = row.get(0);
+    assert_eq!(val, (1,));
+
+    let row = client.query_one("SELECT (1, 'a')", &[]).await.unwrap();
+    let val: (i32, String) = row.get(0);
+    assert_eq!(val, (1, "a".into()));
+
+    let row = client.query_one("SELECT (1, (2, 3))", &[]).await.unwrap();
+    let val: (i32, (i32, i32)) = row.get(0);
+    assert_eq!(val, (1, (2, 3)));
+
+    let row = client.query_one("SELECT (1, 2)", &[]).await.unwrap();
+    let err = row.try_get::<_, (i32, String)>(0).unwrap_err();
+    match err.source() {
+        Some(e) if e.is::<WrongType>() => {}
+        _ => panic!("Unexpected error {:?}", err),
+    };
+
+    let row = client.query_one("SELECT (1, 2, 3)", &[]).await.unwrap();
+    let err = row.try_get::<_, (i32, i32)>(0).unwrap_err();
+    assert_eq!(
+        err.to_string(),
+        "error deserializing column 0: \
+        Postgres record field count does not match Rust tuple length: 3 vs 2"
+    );
+
+    client
+        .batch_execute(
+            "CREATE TYPE pg_temp.simple AS (
+                a int,
+                b text
+            )",
+        )
+        .await
+        .unwrap();
+
+    let row = client
+        .query_one("SELECT (1, 'a')::simple", &[])
+        .await
+        .unwrap();
+    let val: (i32, String) = row.get(0);
+    assert_eq!(val, (1, "a".into()));
+}
+
 #[allow(clippy::eq_op)]
 async fn test_nan_param<T>(sql_type: &str)
 where