Fix panic & query router bug (#85)

* Fix query router bug

* Fix panic

src/client.rs

@@ -302,7 +302,11 @@ impl Client {
             // Handle all custom protocol commands, if any.
             match query_router.try_execute_command(message.clone()) {
                 // Normal query, not a custom command.
-                None => (),
+                None => {
+                    if query_router.query_parser_enabled() {
+                        query_router.infer_role(message.clone());
+                    }
+                }
 
                 // SET SHARD TO
                 Some((Command::SetShard, _)) => {

src/query_router.rs

@@ -126,11 +126,7 @@ impl QueryRouter {
         // This is not a custom query, try to infer which
         // server it'll go to if the query parser is enabled.
         if matches.len() != 1 {
-            debug!("Regular query");
-            if self.query_parser_enabled && self.role() == None {
-                debug!("Inferring role");
-                self.infer_role(buf.clone());
-            }
+            debug!("Regular query, not a command");
             return None;
         }
 
@@ -266,6 +262,8 @@ impl QueryRouter {
 
     /// Try to infer which server to connect to based on the contents of the query.
     pub fn infer_role(&mut self, mut buf: BytesMut) -> bool {
+        debug!("Inferring role");
+
         let code = buf.get_u8() as char;
         let len = buf.get_i32() as usize;