More infrastructure for lazy queries

cc rust-postgres#13

Author: sfackler

src/lib.rs

@@ -482,6 +482,43 @@ impl<'self> NormalPostgresStatement<'self> {
             }
         })
     }
+
+    fn lazy_query<'a>(&'a self, row_limit: uint, params: &[&ToSql])
+            -> PostgresResult<'a> {
+        match self.try_lazy_query(row_limit, params) {
+            Ok(result) => result,
+            Err(err) => fail2!("Error executing query: {}", err.to_str())
+        }
+    }
+
+    fn try_lazy_query<'a>(&'a self, _row_limit: uint, params: &[&ToSql])
+            -> Result<PostgresResult<'a>, PostgresDbError> {
+        match self.execute("", params) {
+            Some(err) => {
+                return Err(err);
+            }
+            None => ()
+        }
+
+        let mut data = ~[];
+        loop {
+            match_read_message_or_fail!(self.conn, {
+                EmptyQueryResponse |
+                CommandComplete {_} => {
+                    break;
+                },
+                DataRow { row } => data.push(row)
+            })
+        }
+        self.conn.wait_for_ready();
+
+        // we're going to be popping off
+        data.reverse();
+        Ok(PostgresResult {
+            stmt: self,
+            data: data,
+        })
+    }
 }
 
 impl<'self> PostgresStatement for NormalPostgresStatement<'self> {
@@ -535,39 +572,12 @@ impl<'self> PostgresStatement for NormalPostgresStatement<'self> {
 
     fn query<'a>(&'a self, params: &[&ToSql])
             -> PostgresResult<'a> {
-        match self.try_query(params) {
-            Ok(result) => result,
-            Err(err) => fail2!("Error running query: {}", err.to_str())
-        }
+        self.lazy_query(0, params)
     }
 
     fn try_query<'a>(&'a self, params: &[&ToSql])
             -> Result<PostgresResult<'a>, PostgresDbError> {
-        match self.execute("", params) {
-            Some(err) => {
-                return Err(err);
-            }
-            None => ()
-        }
-
-        let mut data = ~[];
-        loop {
-            match_read_message_or_fail!(self.conn, {
-                EmptyQueryResponse |
-                CommandComplete {_} => {
-                    break;
-                },
-                DataRow { row } => data.push(row)
-            })
-        }
-        self.conn.wait_for_ready();
-
-        // we're going to be popping off
-        data.reverse();
-        Ok(PostgresResult {
-            stmt: self,
-            data: data,
-        })
+        self.try_lazy_query(0, params)
     }
 
     fn find_col_named(&self, col: &str) -> Option<uint> {
@@ -608,6 +618,18 @@ impl<'self> PostgresStatement for TransactionalPostgresStatement<'self> {
     }
 }
 
+impl<'self> TransactionalPostgresStatement<'self> {
+    pub fn lazy_query<'a>(&'a self, row_limit: uint, params: &[&ToSql])
+            -> PostgresResult<'a> {
+        self.stmt.lazy_query(row_limit, params)
+    }
+
+    pub fn try_lazy_query<'a>(&'a self, row_limit: uint, params: &[&ToSql])
+            -> Result<PostgresResult<'a>, PostgresDbError> {
+        self.try_lazy_query(row_limit, params)
+    }
+}
+
 pub struct PostgresResult<'self> {
     priv stmt: &'self NormalPostgresStatement<'self>,
     priv data: ~[~[Option<~[u8]>]]

src/test.rs

@@ -68,6 +68,26 @@ fn test_query() {
     assert_eq!(~[1i64, 2], result.map(|row| { row[0] }).collect());
 }
 
+#[test]
+fn test_lazy_query() {
+    let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
+
+    do conn.in_transaction |trans| {
+        trans.update("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)", []);
+        let stmt = trans.prepare("INSERT INTO foo (id) VALUES ($1)");
+        let values = ~[0i32, 1, 2, 3, 4, 5];
+        for value in values.iter() {
+            stmt.update([value as &ToSql]);
+        }
+
+        let stmt = trans.prepare("SELECT id FROM foo ORDER BY id");
+        let result = stmt.lazy_query(2, []);
+        assert_eq!(values, result.map(|row| { row[0] }).collect());
+
+        trans.set_rollback();
+    }
+}
+
 fn test_type<T: Eq+ToSql+FromSql>(sql_type: &str, values: &[T]) {
     let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432");
     conn.update("CREATE TEMPORARY TABLE foo (