A 'USE schema' query invokes 'on_init' (#12)

Author: tslilyai

src/lib.rs

@@ -300,8 +300,8 @@ impl<B: MysqlShim<W>, R: Read, W: Write> MysqlIntermediary<B, R, W> {
             let cmd = commands::parse(&packet).unwrap().1;
             match cmd {
                 Command::Query(q) => {
-                    let w = QueryResultWriter::new(&mut self.writer, false);
                     if q.starts_with(b"SELECT @@") || q.starts_with(b"select @@") {
+                        let w = QueryResultWriter::new(&mut self.writer, false);
                         let var = &q[b"SELECT @@".len()..];
                         match var {
                             b"max_allowed_packet" => {
@@ -319,7 +319,16 @@ impl<B: MysqlShim<W>, R: Read, W: Write> MysqlIntermediary<B, R, W> {
                                 w.completed(0, 0)?;
                             }
                         }
+                    } else if q.starts_with(b"USE ") || q.starts_with(b"use ") {
+                        let w = InitWriter {
+                            writer: &mut self.writer,
+                        };
+                        let schema = ::std::str::from_utf8(&q[b"USE ".len()..])
+                            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+                        let schema = schema.trim().trim_end_matches(';').trim_matches('`');
+                        self.shim.on_init(&schema, w)?;
                     } else {
+                        let w = QueryResultWriter::new(&mut self.writer, false);
                         self.shim.on_query(
                             ::std::str::from_utf8(q)
                                 .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?,

tests/main.rs

@@ -155,6 +155,23 @@ fn it_inits_error() {
     .test(|db| assert_eq!(false, db.select_db("test")));
 }
 
+#[test]
+fn it_inits_on_use_query_ok() {
+    TestingShim::new(
+        |_, _| unreachable!(),
+        |_| unreachable!(),
+        |_, _, _| unreachable!(),
+        |schema, writer| {
+            assert_eq!(schema, "test");
+            writer.ok()
+        },
+    )
+    .test(|db| match db.query_drop("USE `test`;") {
+        Ok(_) => assert!(true),
+        Err(_) => assert!(false),
+    });
+}
+
 #[test]
 fn it_pings() {
     TestingShim::new(