Add explicit statement caching

Also remove transaction depth checks to preparation methods on
Connection since lifetimes of statements are tied to the connection, not
any transaction that may be active.

cc rust-postgres#84

Author: sfackler

src/lib.rs

@@ -383,13 +383,21 @@ pub fn cancel_query<T>(params: T, ssl: &SslMode, data: CancelData)
     Ok(())
 }
 
+#[derive(Clone)]
+struct CachedStatement {
+    name: String,
+    param_types: Vec<Type>,
+    result_desc: Vec<ResultDescription>,
+}
+
 struct InnerConnection {
     stream: BufferedStream<MaybeSslStream<InternalStream>>,
     next_stmt_id: usize,
     notice_handler: Box<NoticeHandler>,
     notifications: RingBuf<Notification>,
     cancel_data: CancelData,
     unknown_types: HashMap<Oid, Type>,
+    cached_statements: HashMap<String, CachedStatement>,
     desynchronized: bool,
     finished: bool,
     trans_depth: u32,
@@ -421,6 +429,7 @@ impl InnerConnection {
             notifications: RingBuf::new(),
             cancel_data: CancelData { process_id: 0, secret_key: 0 },
             unknown_types: HashMap::new(),
+            cached_statements: HashMap::new(),
             desynchronized: false,
             finished: false,
             trans_depth: 0,
@@ -647,6 +656,34 @@ impl InnerConnection {
         })
     }
 
+    fn prepare_cached<'a>(&mut self, query: &str, conn: &'a Connection) -> Result<Statement<'a>> {
+        let stmt = self.cached_statements.get(query).map(|e| e.clone());
+
+        let CachedStatement { name, param_types, result_desc } = match stmt {
+            Some(stmt) => stmt,
+            None => {
+                let stmt_name = self.make_stmt_name();
+                let (param_types, result_desc) = try!(self.raw_prepare(&*stmt_name, query));
+                let stmt = CachedStatement {
+                    name: stmt_name,
+                    param_types: param_types,
+                    result_desc: result_desc,
+                };
+                self.cached_statements.insert(query.to_owned(), stmt.clone());
+                stmt
+            }
+        };
+
+        Ok(Statement {
+            conn: conn,
+            name: name,
+            param_types: param_types,
+            result_desc: result_desc,
+            next_portal_id: Cell::new(0),
+            finished: true, // << !
+        })
+    }
+
     fn prepare_copy_in<'a>(&mut self, table: &str, rows: &[&str], conn: &'a Connection)
                            -> Result<CopyInStatement<'a>> {
         let mut query = vec![];
@@ -924,11 +961,31 @@ impl Connection {
     ///     Err(err) => panic!("Error preparing statement: {:?}", err)
     /// };
     pub fn prepare<'a>(&'a self, query: &str) -> Result<Statement<'a>> {
-        let mut conn = self.conn.borrow_mut();
-        if conn.trans_depth != 0 {
-            return Err(Error::WrongTransaction);
-        }
-        conn.prepare(query, self)
+        self.conn.borrow_mut().prepare(query, self)
+    }
+
+    /// Creates cached prepared statement.
+    ///
+    /// Like `prepare`, except that the statement is only prepared once and
+    /// then cached. If the same statement is going to be used frequently,
+    /// caching it can improve performance by reducing the number of round
+    /// trips to the Postgres backend.
+    ///
+    /// ## Example
+    ///
+    /// ```rust,no_run
+    /// # use postgres::{Connection, SslMode};
+    /// # fn f() -> postgres::Result<()> {
+    /// # let x = 10i32;
+    /// # let conn = Connection::connect("", &SslMode::None).unwrap();
+    /// let stmt = try!(conn.prepare_cached("SELECT foo FROM bar WHERE baz = $1"));
+    /// for row in try!(stmt.query(&[&x])) {
+    ///     println!("foo: {}", row.get::<_, String>(0));
+    /// }
+    /// # Ok(()) };
+    /// ```
+    pub fn prepare_cached<'a>(&'a self, query: &str) -> Result<Statement<'a>> {
+        self.conn.borrow_mut().prepare_cached(query, self)
     }
 
     /// Creates a new COPY FROM STDIN prepared statement.
@@ -937,11 +994,7 @@ impl Connection {
     /// the database.
     pub fn prepare_copy_in<'a>(&'a self, table: &str, rows: &[&str])
                                -> Result<CopyInStatement<'a>> {
-        let mut conn = self.conn.borrow_mut();
-        if conn.trans_depth != 0 {
-            return Err(Error::WrongTransaction);
-        }
-        conn.prepare_copy_in(table, rows, self)
+        self.conn.borrow_mut().prepare_copy_in(table, rows, self)
     }
 
     /// Begins a new transaction.
@@ -1039,11 +1092,7 @@ impl Connection {
     /// }
     /// ```
     pub fn batch_execute(&self, query: &str) -> Result<()> {
-        let mut conn = self.conn.borrow_mut();
-        if conn.trans_depth != 0 {
-            return Err(Error::WrongTransaction);
-        }
-        conn.quick_query(query).map(|_| ())
+        self.conn.borrow_mut().quick_query(query).map(|_| ())
     }
 
     /// Returns information used to cancel pending queries.
@@ -1135,12 +1184,20 @@ impl<'conn> Transaction<'conn> {
 
     /// Like `Connection::prepare`.
     pub fn prepare(&self, query: &str) -> Result<Statement<'conn>> {
-        self.conn.conn.borrow_mut().prepare(query, self.conn)
+        self.conn.prepare(query)
+    }
+
+    /// Like `Connection::prepare_cached`.
+    ///
+    /// Note that the statement will be cached for the duration of the
+    /// connection, not just the duration of this transaction.
+    pub fn prepare_cached(&self, query: &str) -> Result<Statement<'conn>> {
+        self.conn.prepare_cached(query)
     }
 
     /// Like `Connection::prepare_copy_in`.
     pub fn prepare_copy_in(&self, table: &str, cols: &[&str]) -> Result<CopyInStatement<'conn>> {
-        self.conn.conn.borrow_mut().prepare_copy_in(table, cols, self.conn)
+        self.conn.prepare_copy_in(table, cols)
     }
 
     /// Like `Connection::execute`.
@@ -1150,7 +1207,7 @@ impl<'conn> Transaction<'conn> {
 
     /// Like `Connection::batch_execute`.
     pub fn batch_execute(&self, query: &str) -> Result<()> {
-        self.conn.conn.borrow_mut().quick_query(query).map(|_| ())
+        self.conn.batch_execute(query)
     }
 
     /// Like `Connection::transaction`.
@@ -1895,6 +1952,9 @@ pub trait GenericConnection {
     /// Like `Connection::prepare`.
     fn prepare<'a>(&'a self, query: &str) -> Result<Statement<'a>>;
 
+    /// Like `Connection::prepare_cached`.
+    fn prepare_cached<'a>(&'a self, query: &str) -> Result<Statement<'a>>;
+
     /// Like `Connection::execute`.
     fn execute(&self, query: &str, params: &[&ToSql]) -> Result<usize>;
 
@@ -1914,6 +1974,10 @@ impl GenericConnection for Connection {
         self.prepare(query)
     }
 
+    fn prepare_cached<'a>(&'a self, query: &str) -> Result<Statement<'a>> {
+        self.prepare_cached(query)
+    }
+
     fn execute(&self, query: &str, params: &[&ToSql]) -> Result<usize> {
         self.execute(query, params)
     }
@@ -1937,6 +2001,10 @@ impl<'a> GenericConnection for Transaction<'a> {
         self.prepare(query)
     }
 
+    fn prepare_cached<'b>(&'b self, query: &str) -> Result<Statement<'b>> {
+        self.prepare_cached(query)
+    }
+
     fn execute(&self, query: &str, params: &[&ToSql]) -> Result<usize> {
         self.execute(query, params)
     }

tests/test.rs

@@ -281,14 +281,9 @@ fn test_nested_transactions_finish() {
 }
 
 #[test]
-fn test_conn_prepare_with_trans() {
+fn test_conn_trans_when_nested() {
     let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
     let _trans = or_panic!(conn.transaction());
-    match conn.prepare("") {
-        Err(Error::WrongTransaction) => {}
-        Err(r) => panic!("Unexpected error {:?}", r),
-        Ok(_) => panic!("Unexpected success"),
-    }
     match conn.transaction() {
         Err(Error::WrongTransaction) => {}
         Err(r) => panic!("Unexpected error {:?}", r),
@@ -902,3 +897,19 @@ fn test_custom_range_element_type() {
         t => panic!("Unexpected type {:?}", t)
     }
 }
+
+#[test]
+fn test_prepare_cached() {
+    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
+    or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT)", &[]));
+    or_panic!(conn.execute("INSERT INTO foo (id) VALUES (1), (2)", &[]));
+
+    let stmt = or_panic!(conn.prepare_cached("SELECT id FROM foo ORDER BY id"));
+    assert_eq!(&[1, 2][], or_panic!(stmt.query(&[])).map(|r| r.get(0)).collect::<Vec<i32>>());
+
+    let stmt = or_panic!(conn.prepare_cached("SELECT id FROM foo ORDER BY id"));
+    assert_eq!(&[1, 2][], or_panic!(stmt.query(&[])).map(|r| r.get(0)).collect::<Vec<i32>>());
+
+    let stmt = or_panic!(conn.prepare_cached("SELECT id FROM foo ORDER BY id DESC"));
+    assert_eq!(&[2, 1][], or_panic!(stmt.query(&[])).map(|r| r.get(0)).collect::<Vec<i32>>());
+}