From 249db6b55a872f1d7fcbbef64908b3ce9bb50531 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Mon, 29 Sep 2014 23:32:57 -0700
Subject: [PATCH] Correctly handle bad column counts in copy

---
 src/lib.rs    | 37 ++++++++++++++++++++++---------------
 tests/test.rs | 27 +++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 15 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index deed27b3..ff5001d4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1614,26 +1614,33 @@ impl<'a> PostgresCopyInStatement<'a> {
         let _ = buf.write_be_i32(0);
         let _ = buf.write_be_i32(0);
 
-        for mut row in rows {
+        for row in rows {
             let _ = buf.write_be_i16(self.column_types.len() as i16);
 
-            let mut count = 0;
-            for (i, (val, ty)) in row.by_ref().zip(self.column_types.iter()).enumerate() {
-                match try!(val.to_sql(ty)) {
-                    (_, None) => {
-                        let _ = buf.write_be_i32(-1);
+            let mut row = row.fuse();
+            let mut types = self.column_types.iter();
+            loop {
+                match (row.next(), types.next()) {
+                    (Some(val), Some(ty)) => {
+                        match try!(val.to_sql(ty)) {
+                            (_, None) => {
+                                let _ = buf.write_be_i32(-1);
+                            }
+                            (_, Some(val)) => {
+                                let _ = buf.write_be_i32(val.len() as i32);
+                                let _ = buf.write(val.as_slice());
+                            }
+                        }
                     }
-                    (_, Some(val)) => {
-                        let _ = buf.write_be_i32(val.len() as i32);
-                        let _ = buf.write(val.as_slice());
+                    (Some(_), None) | (None, Some(_)) => {
+                        try_pg!(conn.stream.write_message(
+                            &CopyFail {
+                                message: "Invalid column count",
+                            }));
+                        break;
                     }
+                    (None, None) => break
                 }
-                count = i+1;
-            }
-
-            if row.next().is_some() || count != self.column_types.len() {
-                // FIXME
-                fail!()
             }
 
             try_pg!(conn.stream.write_message(
diff --git a/tests/test.rs b/tests/test.rs
index 54596f60..c14ce380 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -723,3 +723,30 @@ fn test_copy_in() {
     assert_eq!(vec![(0i32, Some("Steven".to_string())), (1, None)],
                or_fail!(stmt.query([])).map(|r| (r.get(0u), r.get(1u))).collect());
 }
+
+#[test]
+fn test_copy_in_bad_column_count() {
+    let conn = or_fail!(PostgresConnection::connect("postgres://postgres@localhost", &NoSsl));
+    or_fail!(conn.execute("CREATE TEMPORARY TABLE foo (id INT, name VARCHAR)", []));
+
+    let stmt = or_fail!(conn.prepare_copy_in("foo", ["id", "name"]));
+    let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32]];
+
+    let res = stmt.execute(data.iter().map(|r| r.iter().map(|&e| e)));
+    match res {
+        Err(PgDbError(ref err)) if err.message.as_slice().contains("Invalid column count") => {}
+        Err(err) => fail!("unexpected error {}", err),
+        _ => fail!("Expected error"),
+    }
+
+    let data: &[&[&ToSql]] = &[&[&0i32, &"Steven".to_string()], &[&1i32, &"Steven".to_string(), &1i32]];
+
+    let res = stmt.execute(data.iter().map(|r| r.iter().map(|&e| e)));
+    match res {
+        Err(PgDbError(ref err)) if err.message.as_slice().contains("Invalid column count") => {}
+        Err(err) => fail!("unexpected error {}", err),
+        _ => fail!("Expected error"),
+    }
+
+    or_fail!(conn.execute("SELECT 1", []));
+}