Improve error handling in COPY

Author: sfackler

src/lib.rs

@@ -1644,7 +1644,7 @@ impl<'a> PostgresCopyInStatement<'a> {
         let _ = buf.write_be_i32(0);
         let _ = buf.write_be_i32(0);
 
-        for mut row in rows {
+        'l: for mut row in rows {
             let _ = buf.write_be_i16(self.column_types.len() as i16);
 
             let mut types = self.column_types.iter();
@@ -1662,31 +1662,30 @@ impl<'a> PostgresCopyInStatement<'a> {
                         }
                     }
                     (Some(_), None) | (None, Some(_)) => {
-                        try_pg!(conn.stream.write_message(
+                        try_pg_desync!(conn, conn.stream.write_message(
                             &CopyFail {
                                 message: "Invalid column count",
                             }));
-                        break;
+                        break 'l;
                     }
                     (None, None) => break
                 }
             }
 
-            try_pg!(conn.stream.write_message(
+            try_pg_desync!(conn, conn.stream.write_message(
                 &CopyData {
                     data: buf.unwrap().as_slice()
                 }));
             buf = MemWriter::new();
         }
 
         let _ = buf.write_be_i16(-1);
-        try_pg!(conn.stream.write_message(
-            &CopyData {
+        try_pg!(conn.write_messages([
+            CopyData {
                 data: buf.unwrap().as_slice(),
-            }));
-        try_pg!(conn.stream.write_message(&CopyDone));
-        try_pg!(conn.stream.write_message(&Sync));
-        try_pg!(conn.stream.flush());
+            },
+            CopyDone,
+            Sync]));
 
         let num = match try_pg!(conn.read_message_()) {
             CommandComplete { tag } => util::parse_update_count(tag),

src/macros.rs

@@ -28,6 +28,18 @@ macro_rules! try_desync(
     )
 )
 
+macro_rules! try_pg_desync(
+    ($s:expr, $e:expr) => (
+        match $e {
+            Ok(ok) => ok,
+            Err(err) => {
+                $s.desynchronized = true;
+                return Err(PgStreamError(err))
+            }
+        }
+    )
+)
+
 macro_rules! check_desync(
     ($e:expr) => ({
         if $e.canary() != CANARY {