Enable returning errors after creating row writer (jonhoo#20)

Authored-by: Simon Dickson <[email protected]>

Author: simonhdickson

src/resultset.rs

@@ -347,27 +347,31 @@ where
 }
 
 impl<'a, W: Write + 'a> RowWriter<'a, W> {
-    fn finish_inner(&mut self) -> io::Result<()> {
+    fn finish_inner(&mut self, complete: bool) -> io::Result<()> {
         if self.finished {
             return Ok(());
         }
+
         self.finished = true;
 
         if !self.columns.is_empty() && self.col != 0 {
             self.end_row()?;
         }
 
-        if self.columns.is_empty() {
-            // response to no column query is always an OK packet
-            // we've kept track of the number of rows in col (hacky, I know)
-            self.result.as_mut().unwrap().last_end = Some(Finalizer::Ok {
-                rows: self.col as u64,
-                last_insert_id: 0,
-            });
-        } else {
-            // we wrote out at least one row
-            self.result.as_mut().unwrap().last_end = Some(Finalizer::EOF);
+        if complete {
+            if self.columns.is_empty() {
+                // response to no column query is always an OK packet
+                // we've kept track of the number of rows in col (hacky, I know)
+                self.result.as_mut().unwrap().last_end = Some(Finalizer::Ok {
+                    rows: self.col as u64,
+                    last_insert_id: 0,
+                });
+            } else {
+                // we wrote out at least one row
+                self.result.as_mut().unwrap().last_end = Some(Finalizer::EOF);
+            }
         }
+
         Ok(())
     }
 
@@ -378,15 +382,26 @@ impl<'a, W: Write + 'a> RowWriter<'a, W> {
 
     /// End this resultset response, and indicate to the client that no more rows are coming.
     pub fn finish_one(mut self) -> io::Result<QueryResultWriter<'a, W>> {
-        self.finish_inner()?;
+        self.finish_inner(true)?;
+
         // we know that dropping self will see self.finished == true,
         // and so Drop won't try to use self.result.
         Ok(self.result.take().unwrap())
     }
+
+    /// End this resultset response, and indicate to the client there was an error.
+    pub fn finish_error<E>(mut self, kind: ErrorKind, msg: &E) -> io::Result<()>
+    where
+        E: Borrow<[u8]>,
+    {
+        self.finish_inner(false)?;
+
+        self.result.take().unwrap().error(kind, msg)
+    }
 }
 
 impl<'a, W: Write + 'a> Drop for RowWriter<'a, W> {
     fn drop(&mut self) {
-        self.finish_inner().unwrap();
+        self.finish_inner(true).unwrap();
     }
 }

tests/main.rs

@@ -266,6 +266,44 @@ fn error_response() {
     })
 }
 
+#[test]
+fn error_in_result_set_response() {
+    let err = (ErrorKind::ER_NO, "clearly not");
+    TestingShim::new(
+        move |_, w| {
+            let cols = &[Column {
+                table: String::new(),
+                column: "a".to_owned(),
+                coltype: myc::constants::ColumnType::MYSQL_TYPE_SHORT,
+                colflags: myc::constants::ColumnFlags::empty(),
+            }];
+            let mut w = w.start(cols)?;
+            w.write_col(1024)?;
+            w.finish_error(err.0, &err.1.as_bytes())
+        },
+        |_| unreachable!(),
+        |_, _, _| unreachable!(),
+        |_, _| unreachable!(),
+    )
+    .test(|db| {
+        let mut result = db.query_iter("SELECT a FROM foo").unwrap();
+        let row1 = result.next().unwrap().unwrap().get::<i16, _>(0).unwrap();
+        assert_eq!(row1, 1024);
+        if let mysql::Error::MySqlError(e) = result.by_ref().next().unwrap().unwrap_err() {
+            assert_eq!(
+                e,
+                mysql::error::MySqlError {
+                    state: String::from_utf8(err.0.sqlstate().to_vec()).unwrap(),
+                    message: err.1.to_owned(),
+                    code: err.0 as u16,
+                }
+            );
+        } else {
+            unreachable!()
+        }
+    })
+}
+
 #[test]
 fn empty_on_drop() {
     let cols = [Column {