From 86941d62e47dc4a3d75ea49c5f13172fa99cd4f2 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <levkk@users.noreply.github.com>
Date: Mon, 21 Feb 2022 00:00:50 -0800
Subject: [PATCH] Reset query router setting to default (#32)

---
 .circleci/config.yml |  3 +++
 src/pool.rs          |  5 +----
 src/query_router.rs  | 14 ++++++++------
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index 6b330fe..63bfee5 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -25,6 +25,9 @@ jobs:
       - checkout
       - restore_cache:
           key: cargo-lock-2-{{ checksum "Cargo.lock" }}
+      - run:
+          name: "Lint"
+          command: "cargo fmt --check"
       - run:
           name: "Install dependencies"
           command: "sudo apt-get update && sudo apt-get install -y psmisc postgresql-contrib-12 postgresql-client-12 ruby ruby-dev libpq-dev"
diff --git a/src/pool.rs b/src/pool.rs
index 87f169a..6c9a7a2 100644
--- a/src/pool.rs
+++ b/src/pool.rs
@@ -167,10 +167,7 @@ impl ConnectionPool {
         };
 
         if !exists {
-            error!(
-                "Requested role {:?}, but none are configured",
-                role
-            );
+            error!("Requested role {:?}, but none are configured", role);
             return Err(Error::BadConfig);
         }
 
diff --git a/src/query_router.rs b/src/query_router.rs
index 64d6e70..4134af8 100644
--- a/src/query_router.rs
+++ b/src/query_router.rs
@@ -3,12 +3,12 @@ use crate::sharding::Sharder;
 /// Route queries automatically based on explicitely requested
 /// or implied query characteristics.
 use bytes::{Buf, BytesMut};
+use log::{debug, error};
 use once_cell::sync::OnceCell;
 use regex::RegexSet;
 use sqlparser::ast::Statement::{Query, StartTransaction};
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
-use log::{error, debug};
 
 const CUSTOM_SQL_REGEXES: [&str; 5] = [
     r"(?i)SET SHARDING KEY TO '[0-9]+'",
@@ -174,6 +174,7 @@ impl QueryRouter {
                     "default" => {
                         // TODO: reset query parser to default here.
                         self.active_role = self.default_server_role;
+                        self.query_parser_enabled = get_config().query_router.query_parser_enabled;
                         self.active_role
                     }
 
@@ -220,11 +221,7 @@ impl QueryRouter {
         let ast = match Parser::parse_sql(&PostgreSqlDialect {}, &query) {
             Ok(ast) => ast,
             Err(err) => {
-                debug!(
-                    "{:?}, query: {}",
-                    err,
-                    query
-                );
+                debug!("{:?}, query: {}", err, query);
                 return false;
             }
         };
@@ -483,5 +480,10 @@ mod test {
         let query = simple_query("SELECT * FROM test_table");
         assert_eq!(qr.infer_role(query), true);
         assert_eq!(qr.role(), Some(Role::Replica));
+
+        assert!(qr.query_parser_enabled());
+        let query = simple_query("SET SERVER ROLE TO 'default'");
+        assert!(qr.try_execute_command(query) != None);
+        assert!(!qr.query_parser_enabled());
     }
 }