Disallow use of wrong transactions

We don't want to allow statements to outlive the active transaction in
which they were prepared.

Author: sfackler

src/error.rs

@@ -535,6 +535,9 @@ pub enum PostgresError {
     PgInvalidColumn,
     /// A value was NULL but converted to a non-nullable Rust type
     PgWasNull,
+    /// An attempt was made to prepare a statement or start a transaction on an
+    /// object other than the active transaction
+    PgWrongTransaction,
 }
 
 impl fmt::Show for PostgresError {
@@ -554,6 +557,10 @@ impl fmt::Show for PostgresError {
             PgWrongType(ref ty) => write!(fmt, "Unexpected type {}", ty),
             PgInvalidColumn => write!(fmt, "Invalid column"),
             PgWasNull => write!(fmt, "The value was NULL"),
+            PgWrongTransaction =>
+                write!(fmt, "An attempt was made to prepare a statement or \
+                             start a transaction on an object other than the \
+                             active transaction"),
         }
     }
 }

src/lib.rs

@@ -105,7 +105,8 @@ use error::{InvalidUrl,
             PostgresDbError,
             PostgresError,
             UnsupportedAuthentication,
-            PgWrongConnection};
+            PgWrongConnection,
+            PgWrongTransaction};
 use io::{MaybeSslStream,
          InternalStream};
 use message::{AuthenticationCleartextPassword,
@@ -415,6 +416,7 @@ struct InnerPostgresConnection {
     unknown_types: HashMap<Oid, StrBuf>,
     desynchronized: bool,
     finished: bool,
+    trans_depth: u32,
     canary: u32,
 }
 
@@ -442,6 +444,7 @@ impl InnerPostgresConnection {
             unknown_types: HashMap::new(),
             desynchronized: false,
             finished: false,
+            trans_depth: 0,
             canary: CANARY,
         };
 
@@ -805,7 +808,11 @@ impl PostgresConnection {
     /// };
     pub fn prepare<'a>(&'a self, query: &str)
             -> PostgresResult<PostgresStatement<'a>> {
-        self.conn.borrow_mut().prepare(query, self)
+        let mut conn = self.conn.borrow_mut();
+        if conn.trans_depth != 0 {
+            return Err(PgWrongTransaction);
+        }
+        conn.prepare(query, self)
     }
 
     /// Begins a new transaction.
@@ -837,11 +844,15 @@ impl PostgresConnection {
     pub fn transaction<'a>(&'a self)
             -> PostgresResult<PostgresTransaction<'a>> {
         check_desync!(self);
+        if self.conn.borrow().trans_depth != 0 {
+            return Err(PgWrongTransaction);
+        }
         try!(self.quick_query("BEGIN"));
+        self.conn.borrow_mut().trans_depth += 1;
         Ok(PostgresTransaction {
             conn: self,
             commit: Cell::new(true),
-            nested: false,
+            depth: 1,
             finished: false,
         })
     }
@@ -921,7 +932,7 @@ pub enum SslMode {
 pub struct PostgresTransaction<'conn> {
     conn: &'conn PostgresConnection,
     commit: Cell<bool>,
-    nested: bool,
+    depth: u32,
     finished: bool,
 }
 
@@ -936,51 +947,50 @@ impl<'conn> Drop for PostgresTransaction<'conn> {
 
 impl<'conn> PostgresTransaction<'conn> {
     fn finish_inner(&mut self) -> PostgresResult<()> {
+        debug_assert!(self.depth == self.conn.conn.borrow().trans_depth);
         let rollback = task::failing() || !self.commit.get();
-        let query = match (rollback, self.nested) {
+        let query = match (rollback, self.depth != 1) {
             (true, true) => "ROLLBACK TO sp",
             (true, false) => "ROLLBACK",
             (false, true) => "RELEASE sp",
             (false, false) => "COMMIT",
         };
+        self.conn.conn.borrow_mut().trans_depth -= 1;
         self.conn.quick_query(query).map(|_| ())
     }
 
     /// Like `PostgresConnection::prepare`.
     pub fn prepare<'a>(&'a self, query: &str)
             -> PostgresResult<PostgresStatement<'a>> {
-        self.conn.prepare(query)
+        if self.conn.conn.borrow().trans_depth != self.depth {
+            return Err(PgWrongTransaction);
+        }
+        self.conn.conn.borrow_mut().prepare(query, self.conn)
     }
 
     /// Like `PostgresConnection::execute`.
     pub fn execute(&self, query: &str, params: &[&ToSql])
             -> PostgresResult<uint> {
-        self.conn.execute(query, params)
+        self.prepare(query).and_then(|s| s.execute(params))
     }
 
     /// Like `PostgresConnection::transaction`.
     pub fn transaction<'a>(&'a self)
             -> PostgresResult<PostgresTransaction<'a>> {
         check_desync!(self.conn);
+        if self.conn.conn.borrow().trans_depth != self.depth {
+            return Err(PgWrongTransaction);
+        }
         try!(self.conn.quick_query("SAVEPOINT sp"));
+        self.conn.conn.borrow_mut().trans_depth += 1;
         Ok(PostgresTransaction {
             conn: self.conn,
             commit: Cell::new(true),
-            nested: true,
+            depth: self.depth + 1,
             finished: false,
         })
     }
 
-    /// Like `PostgresConnection::notifications`.
-    pub fn notifications<'a>(&'a self) -> PostgresNotifications<'a> {
-        self.conn.notifications()
-    }
-
-    /// Like `PostgresConnection::is_desynchronized`.
-    pub fn is_desynchronized(&self) -> bool {
-        self.conn.is_desynchronized()
-    }
-
     /// Determines if the transaction is currently set to commit or roll back.
     pub fn will_commit(&self) -> bool {
         self.commit.get()
@@ -1447,12 +1457,8 @@ impl RowIndex for int {
 impl<'a> RowIndex for &'a str {
     #[inline]
     fn idx(&self, stmt: &PostgresStatement) -> Option<uint> {
-        for (i, desc) in stmt.result_descriptions().iter().enumerate() {
-            if desc.name.as_slice() == *self {
-                return Some(i);
-            }
-        }
-        None
+        stmt.result_descriptions().iter()
+            .position(|d| d.name.as_slice() == *self)
     }
 }
 

src/test.rs

@@ -225,7 +225,7 @@ fn test_nested_transactions() {
             }
         }
 
-        let stmt = or_fail!(conn.prepare("SELECT * FROM foo ORDER BY id"));
+        let stmt = or_fail!(trans1.prepare("SELECT * FROM foo ORDER BY id"));
         let result = or_fail!(stmt.query([]));
 
         assert_eq!(vec![1i32, 2, 4, 6], result.map(|row| row[1]).collect());
@@ -277,10 +277,13 @@ fn test_nested_transactions_finish() {
             assert!(trans2.finish().is_ok());
         }
 
-        let stmt = or_fail!(conn.prepare("SELECT * FROM foo ORDER BY id"));
-        let result = or_fail!(stmt.query([]));
+        // in a block to unborrow trans1 for the finish call
+        {
+            let stmt = or_fail!(trans1.prepare("SELECT * FROM foo ORDER BY id"));
+            let result = or_fail!(stmt.query([]));
 
-        assert_eq!(vec![1i32, 2, 4, 6], result.map(|row| row[1]).collect());
+            assert_eq!(vec![1i32, 2, 4, 6], result.map(|row| row[1]).collect());
+        }
 
         trans1.set_rollback();
         assert!(trans1.finish().is_ok());
@@ -332,7 +335,7 @@ fn test_lazy_query() {
     for value in values.iter() {
         or_fail!(stmt.execute([value as &ToSql]));
     }
-    let stmt = or_fail!(conn.prepare("SELECT id FROM foo ORDER BY id"));
+    let stmt = or_fail!(trans.prepare("SELECT id FROM foo ORDER BY id"));
     let result = or_fail!(trans.lazy_query(&stmt, [], 2));
     assert_eq!(values, result.map(|row| row.unwrap()[1]).collect());