From f06f64119c1a3129ce9501ee625518708bdadb67 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Fri, 24 Jun 2022 15:14:31 -0700
Subject: [PATCH] Fix panic & query router bug (#85)

* Fix query router bug

* Fix panic
---
 src/client.rs       | 6 +++++-
 src/query_router.rs | 8 +++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index b53bd33..ea1bf01 100644
--- a/src/client.rs
+++ b/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, _)) => {
diff --git a/src/query_router.rs b/src/query_router.rs
index 995a25c..98a4770 100644
--- a/src/query_router.rs
+++ b/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;