Handle errors when running copy statements

Author: sfackler

src/stmt.rs

@@ -310,6 +310,10 @@ impl<'conn> Statement<'conn> {
 
         let (format, column_formats) = match try!(conn.read_message()) {
             CopyInResponse { format, column_formats } => (format, column_formats),
+            ErrorResponse { fields } => {
+                try!(conn.wait_for_ready());
+                return DbError::new(fields);
+            }
             _ => {
                 loop {
                     match try!(conn.read_message()) {
@@ -425,6 +429,10 @@ impl<'conn> Statement<'conn> {
                             io::ErrorKind::InvalidInput,
                             "called `copy_out` on a non-`COPY TO STDOUT` statement")));
             }
+            ErrorResponse { fields } => {
+                try!(conn.wait_for_ready());
+                return DbError::new(fields);
+            }
             _ => {
                 loop {
                     match try!(conn.read_message()) {

tests/test.rs

@@ -811,6 +811,21 @@ fn test_copy_out() {
     or_panic!(conn.batch_execute("SELECT 1"));
 }
 
+#[test]
+fn test_copy_out_error() {
+    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
+    or_panic!(conn.batch_execute("
+         CREATE TEMPORARY TABLE foo (id INT);
+         INSERT INTO foo (id) VALUES (0), (1), (2), (3)"));
+    let stmt = or_panic!(conn.prepare("COPY (SELECT id FROM foo ORDER BY id) TO STDOUT (OIDS)"));
+    let mut buf = vec![];
+    match stmt.copy_out(&[], &mut buf) {
+        Ok(_) => panic!("unexpected success"),
+        Err(Error::DbError(..)) => {}
+        Err(e) => panic!("unexpected error {}", e),
+    }
+}
+
 #[test]
 // Just make sure the impls don't infinite loop
 fn test_generic_connection() {