Properly return ioerrors to the caller in copy_in

Author: sfackler

src/lib.rs

@@ -1581,18 +1581,27 @@ impl<'conn> Statement<'conn> {
                 Ok(0) => break,
                 Ok(len) => {
                     try_desync!(conn, conn.stream.write_message(
-                            &CopyData {
-                                data: &buf[..len as usize],
-                            }));
+                        &CopyData {
+                            data: &buf[..len as usize],
+                        }));
                     buf.clear();
                 }
                 Err(err) => {
-                    // FIXME better to return the error directly
-                    try_desync!(conn, conn.stream.write_message(
-                            &CopyFail {
-                                message: &err.to_string(),
-                            }));
-                    break;
+                    try!(conn.write_messages(&[
+                        CopyFail {
+                            message: "",
+                        },
+                        CopyDone,
+                        Sync]));
+                    match try!(conn.read_message()) {
+                        ErrorResponse { .. } => { /* expected from the CopyFail */ }
+                        _ => {
+                            conn.desynchronized = true;
+                            return Err(Error::IoError(bad_response()));
+                        }
+                    }
+                    try!(conn.wait_for_ready());
+                    return Err(Error::IoError(err));
                 }
             }
         }

tests/test.rs

@@ -7,6 +7,7 @@ extern crate openssl;
 #[cfg(feature = "openssl")]
 use openssl::ssl::{SslContext, SslMethod};
 use std::thread;
+use std::io;
 
 use postgres::{HandleNotice,
                Notification,
@@ -605,48 +606,6 @@ fn test_notifications_next_block() {
     }, or_panic!(notifications.next_block()));
 }
 
-/*
-#[test]
-fn test_notifications_next_block_for() {
-    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
-    or_panic!(conn.execute("LISTEN test_notifications_next_block_for", &[]));
-
-    let _t = thread::spawn(|| {
-        let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
-        timer::sleep(Duration::milliseconds(500));
-        or_panic!(conn.execute("NOTIFY test_notifications_next_block_for, 'foo'", &[]));
-    });
-
-    let mut notifications = conn.notifications();
-    check_notification(Notification {
-        pid: 0,
-        channel: "test_notifications_next_block_for".to_string(),
-        payload: "foo".to_string()
-    }, or_panic!(notifications.next_block_for(Duration::seconds(2)).unwrap()));
-}
-
-#[test]
-fn test_notifications_next_block_for_timeout() {
-    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
-    or_panic!(conn.execute("LISTEN test_notifications_next_block_for_timeout", &[]));
-
-    let _t = thread::spawn(|| {
-        let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
-        timer::sleep(Duration::seconds(2));
-        or_panic!(conn.execute("NOTIFY test_notifications_next_block_for_timeout, 'foo'", &[]));
-    });
-
-    let mut notifications = conn.notifications();
-    match notifications.next_block_for(Duration::milliseconds(500)) {
-        None => {}
-        Some(Err(e)) => panic!("Unexpected error {:?}", e),
-        Some(Ok(_)) => panic!("expected error"),
-    }
-
-    or_panic!(conn.execute("SELECT 1", &[]));
-}
-*/
-
 #[test]
 // This test is pretty sad, but I don't think there's a better way :(
 fn test_cancel_query() {
@@ -762,6 +721,28 @@ fn test_batch_execute_copy_from_err() {
     }
 }
 
+#[test]
+fn test_copy_io_error() {
+    struct ErrorReader;
+
+    impl io::Read for ErrorReader {
+        fn read(&mut self, _: &mut [u8]) -> io::Result<usize> {
+            Err(io::Error::new(io::ErrorKind::AddrNotAvailable, "boom"))
+        }
+    }
+
+    let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));
+    or_panic!(conn.execute("CREATE TEMPORARY TABLE foo (id INT)", &[]));
+    let stmt = or_panic!(conn.prepare("COPY foo (id) FROM STDIN"));
+    match stmt.copy_in(&[], &mut ErrorReader) {
+        Err(Error::IoError(ref e)) if e.kind() == io::ErrorKind::AddrNotAvailable => {}
+        Err(err) => panic!("Unexpected error {:?}", err),
+        _ => panic!("Expected error"),
+    }
+
+    or_panic!(conn.execute("SELECT 1", &[]));
+}
+
 #[test]
 fn test_copy() {
     let conn = or_panic!(Connection::connect("postgres://postgres@localhost", &SslMode::None));