diff --git a/src/lib.rs b/src/lib.rs index 27f297ab..8f2fd0ec 100644 --- a/src/lib.rs +++ b/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, 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, 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 { @@ -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, PostgresDbError> { + self.try_lazy_query(row_limit, params) + } +} + pub struct PostgresResult<'self> { priv stmt: &'self NormalPostgresStatement<'self>, priv data: ~[~[Option<~[u8]>]] diff --git a/src/test.rs b/src/test.rs index 4879d889..2a0c2dfe 100644 --- a/src/test.rs +++ b/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(sql_type: &str, values: &[T]) { let conn = PostgresConnection::connect("postgres://postgres@127.0.0.1:5432"); conn.update("CREATE TEMPORARY TABLE foo (