Redesign COPY IN execution setup.

The old Iterator of Iterators setup basically required all of the rows
to sit in memory at the same time, which defeats the entire purpose of
that configuration. There's probably a library that defines
`StreamIterator`, so this may jump to that.
This commit is contained in:
Steven Fackler 2015-01-21 20:25:01 -08:00
parent 9e260fdfac
commit 290d572166
2 changed files with 75 additions and 32 deletions

View File

@ -906,25 +906,6 @@ impl Connection {
/// ///
/// These statements provide a method to efficiently bulk-upload data to /// These statements provide a method to efficiently bulk-upload data to
/// the database. /// the database.
///
/// ## Example
///
/// ```rust,no_run
/// # use postgres::{Connection, SslMode, Error};
/// # use postgres::types::ToSql;
/// # fn f() -> Result<(), Error> {
/// # let conn = Connection::connect("", &SslMode::None).unwrap();
/// try!(conn.execute("CREATE TABLE foo (
/// bar INT PRIMARY KEY,
/// baz VARCHAR
/// )", &[]));
///
/// let stmt = try!(conn.prepare_copy_in("foo", &["bar", "baz"]));
/// let data: &[&[&ToSql]] = &[&[&0i32, &"blah"],
/// &[&1i32, &None::<String>]];
/// try!(stmt.execute(data.iter().map(|r| r.iter().map(|&e| e))));
/// # Ok(()) };
/// ```
pub fn prepare_copy_in<'a>(&'a self, table: &str, rows: &[&str]) pub fn prepare_copy_in<'a>(&'a self, table: &str, rows: &[&str])
-> Result<CopyInStatement<'a>> { -> Result<CopyInStatement<'a>> {
let mut conn = self.conn.borrow_mut(); let mut conn = self.conn.borrow_mut();
@ -1700,6 +1681,45 @@ impl<'a> Drop for CopyInStatement<'a> {
} }
} }
/// An `Iterator` variant which returns borrowed values.
pub trait StreamIterator {
/// Returns the next value, or `None` if there is none.
fn next<'a>(&'a mut self) -> Option<&'a (ToSql + 'a)>;
}
/// An adapter type implementing `StreamIterator` for a `Vec<Box<ToSql>>`.
pub struct VecStreamIterator<'a> {
v: Vec<Box<ToSql + 'a>>,
idx: usize,
}
impl<'a> VecStreamIterator<'a> {
/// Creates a new `VecStreamIterator`.
pub fn new(v: Vec<Box<ToSql + 'a>>) -> VecStreamIterator<'a> {
VecStreamIterator {
v: v,
idx: 0,
}
}
/// Returns the underlying `Vec`.
pub fn into_inner(self) -> Vec<Box<ToSql + 'a>> {
self.v
}
}
impl<'a> StreamIterator for VecStreamIterator<'a> {
fn next<'b>(&'b mut self) -> Option<&'b (ToSql + 'b)> {
match self.v.get_mut(self.idx) {
Some(mut e) => {
self.idx += 1;
Some(&mut **e)
},
None => None,
}
}
}
impl<'a> CopyInStatement<'a> { impl<'a> CopyInStatement<'a> {
fn finish_inner(&mut self) -> Result<()> { fn finish_inner(&mut self) -> Result<()> {
let mut conn = self.conn.conn.borrow_mut(); let mut conn = self.conn.conn.borrow_mut();
@ -1714,12 +1734,14 @@ impl<'a> CopyInStatement<'a> {
/// Executes the prepared statement. /// Executes the prepared statement.
/// ///
/// Each iterator returned by the `rows` iterator will be interpreted as /// The `rows` argument is an `Iterator` returning `StreamIterator` values,
/// providing a single result row. /// each one of which provides values for a row of input. This setup is
/// designed to allow callers to avoid having to maintain the entire row
/// set in memory.
/// ///
/// Returns the number of rows copied. /// Returns the number of rows copied.
pub fn execute<'b, I, J>(&self, mut rows: I) -> Result<usize> pub fn execute<'b, I, J>(&self, mut rows: I) -> Result<usize>
where I: Iterator<Item=J>, J: Iterator<Item=&'b (ToSql + 'b)> { where I: Iterator<Item=J>, J: StreamIterator {
let mut conn = self.conn.conn.borrow_mut(); let mut conn = self.conn.conn.borrow_mut();
try!(conn.write_messages(&[ try!(conn.write_messages(&[

View File

@ -23,7 +23,8 @@ use postgres::{NoticeHandler,
ToSql, ToSql,
Error, Error,
ConnectError, ConnectError,
DbError}; DbError,
VecStreamIterator};
use postgres::SqlState::{SyntaxError, use postgres::SqlState::{SyntaxError,
QueryCanceled, QueryCanceled,
UndefinedTable, UndefinedTable,
@ -786,12 +787,16 @@ fn test_copy_in() {
or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[])); or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[]));
let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"])); let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"]));
let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32, &None::<String>]];
assert_eq!(Ok(2), stmt.execute(data.iter().map(|r| r.iter().map(|&e| e)))); let data = (0i32..2).map(|i| {
VecStreamIterator::new(vec![Box::new(i) as Box<ToSql>,
Box::new(format!("{}", i)) as Box<ToSql>])
});
assert_eq!(Ok(2), stmt.execute(data));
let stmt = or_panic!(conn.prepare("SELECT id, name FROM foo ORDER BY id")); let stmt = or_panic!(conn.prepare("SELECT id, name FROM foo ORDER BY id"));
assert_eq!(vec![(0i32, Some("Steven".to_string())), (1, None)], assert_eq!(vec![(0i32, Some("0".to_string())), (1, Some("1".to_string()))],
or_panic!(stmt.query(&[])).map(|r| (r.get(0), r.get(1))).collect::<Vec<_>>()); or_panic!(stmt.query(&[])).map(|r| (r.get(0), r.get(1))).collect::<Vec<_>>());
} }
@ -801,18 +806,28 @@ fn test_copy_in_bad_column_count() {
or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[])); or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[]));
let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"])); let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"]));
let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32]]; let data = vec![
VecStreamIterator::new(vec![Box::new(1i32) as Box<ToSql>,
Box::new("Steven".to_string()) as Box<ToSql>]),
VecStreamIterator::new(vec![Box::new(2i32) as Box<ToSql>]),
].into_iter();
let res = stmt.execute(data.iter().map(|r| r.iter().map(|&e| e))); let res = stmt.execute(data);
match res { match res {
Err(Error::DbError(ref err)) if err.message[].contains("Invalid column count") => {} Err(Error::DbError(ref err)) if err.message[].contains("Invalid column count") => {}
Err(err) => panic!("unexpected error {:?}", err), Err(err) => panic!("unexpected error {:?}", err),
_ => panic!("Expected error"), _ => panic!("Expected error"),
} }
let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32, &"Steven".to_string(), &1i32]]; let data = vec![
VecStreamIterator::new(vec![Box::new(1i32) as Box<ToSql>,
Box::new("Steven".to_string()) as Box<ToSql>]),
VecStreamIterator::new(vec![Box::new(2i32) as Box<ToSql>,
Box::new("Steven".to_string()) as Box<ToSql>,
Box::new(3i64) as Box<ToSql>]),
].into_iter();
let res = stmt.execute(data.iter().map(|r| r.iter().map(|&e| e))); let res = stmt.execute(data);
match res { match res {
Err(Error::DbError(ref err)) if err.message[].contains("Invalid column count") => {} Err(Error::DbError(ref err)) if err.message[].contains("Invalid column count") => {}
Err(err) => panic!("unexpected error {:?}", err), Err(err) => panic!("unexpected error {:?}", err),
@ -828,9 +843,15 @@ fn test_copy_in_bad_type() {
or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[])); or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", &[]));
let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"])); let stmt = or_panic!(conn.prepare_copy_in("foo", &["id", "name"]));
let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32, &2i32]];
let res = stmt.execute(data.iter().map(|r| r.iter().map(|&e| e))); let data = vec![
VecStreamIterator::new(vec![Box::new(1i32) as Box<ToSql>,
Box::new("Steven".to_string()) as Box<ToSql>]),
VecStreamIterator::new(vec![Box::new(2i32) as Box<ToSql>,
Box::new(1i32) as Box<ToSql>]),
].into_iter();
let res = stmt.execute(data);
match res { match res {
Err(Error::DbError(ref err)) if err.message[].contains("Unexpected type Varchar") => {} Err(Error::DbError(ref err)) if err.message[].contains("Unexpected type Varchar") => {}
Err(err) => panic!("unexpected error {:?}", err), Err(err) => panic!("unexpected error {:?}", err),