From dc2f04de96b453c900edeb6aa4a6c313b6997c91 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Sat, 15 Aug 2015 20:39:13 -0700
Subject: [PATCH] Start work on COPY TO statements

---
 src/lib.rs     | 25 +++++++++++++++++++++++++
 src/message.rs | 28 ++++++++++++++++++++++++++++
 tests/test.rs  | 13 +++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index e8c25f08..985e1964 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1450,6 +1450,20 @@ impl<'conn> Statement<'conn> {
                         },
                         Sync]));
                 }
+                CopyOutResponse { .. } => {
+                    loop {
+                        match try!(conn.read_message()) {
+                            BCopyDone => break,
+                            ErrorResponse { fields } => {
+                                try!(conn.wait_for_ready());
+                                return DbError::new(fields);
+                            }
+                            _ => {}
+                        }
+                    }
+                    num = 0;
+                    break;
+                }
                 _ => {
                     conn.desynchronized = true;
                     return Err(Error::IoError(bad_response()));
@@ -1690,6 +1704,17 @@ fn read_rows(conn: &mut InnerConnection, buf: &mut VecDeque<Vec<Option<Vec<u8>>>
                     },
                     Sync]));
             }
+            CopyOutResponse { .. } => {
+                loop {
+                    match try!(conn.read_message()) {
+                        ReadyForQuery { .. } => break,
+                        _ => {}
+                    }
+                }
+                return Err(Error::IoError(std_io::Error::new(
+                            std_io::ErrorKind::InvalidInput,
+                            "COPY queries cannot be directly executed")));
+            }
             _ => {
                 conn.desynchronized = true;
                 return Err(Error::IoError(bad_response()));
diff --git a/src/message.rs b/src/message.rs
index 71800df0..765d56d3 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -32,10 +32,19 @@ pub enum BackendMessage {
     CommandComplete {
         tag: String,
     },
+    // FIXME naming
+    BCopyData {
+        data: Vec<u8>,
+    },
+    BCopyDone,
     CopyInResponse {
         format: u8,
         column_formats: Vec<u16>,
     },
+    CopyOutResponse {
+        format: u8,
+        column_formats: Vec<u16>,
+    },
     DataRow {
         row: Vec<Option<Vec<u8>>>
     },
@@ -292,7 +301,15 @@ impl<R: BufRead> ReadMessage for R {
                 channel: try!(rdr.read_cstr()),
                 payload: try!(rdr.read_cstr())
             },
+            b'c' => BCopyDone,
             b'C' => CommandComplete { tag: try!(rdr.read_cstr()) },
+            b'd' => {
+                let mut data = vec![];
+                try!(rdr.read_to_end(&mut data));
+                BCopyData {
+                    data: data,
+                }
+            }
             b'D' => try!(read_data_row(&mut rdr)),
             b'E' => ErrorResponse { fields: try!(read_fields(&mut rdr)) },
             b'G' => {
@@ -306,6 +323,17 @@ impl<R: BufRead> ReadMessage for R {
                     column_formats: column_formats,
                 }
             }
+            b'H' => {
+                let format = try!(rdr.read_u8());
+                let mut column_formats = vec![];
+                for _ in 0..try!(rdr.read_u16::<BigEndian>()) {
+                    column_formats.push(try!(rdr.read_u16::<BigEndian>()));
+                }
+                CopyOutResponse {
+                    format: format,
+                    column_formats: column_formats,
+                }
+            }
             b'I' => EmptyQueryResponse,
             b'K' => BackendKeyData {
                 process_id: try!(rdr.read_u32::<BigEndian>()),
diff --git a/tests/test.rs b/tests/test.rs
index 809bcbaa..d5d8d7b9 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -756,6 +756,19 @@ fn test_copy() {
                stmt.query(&[]).unwrap().iter().map(|r| r.get(0)).collect::<Vec<i32>>());
 }
 
+#[test]
+fn test_copy_out_query() {
+    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 foo (id) TO STDOUT"));
+    match stmt.query(&[]) {
+        Ok(_) => panic!("unexpected success"),
+        Err(Error::IoError(ref e)) if e.to_string().contains("COPY") => {}
+        Err(e) => panic!("unexpected error {:?}", e),
+    }
+}
 
 #[test]
 // Just make sure the impls don't infinite loop