Query parser 3.0 (#23)

* Starting query parsing

* Query parser

* working config

* disable by default

* fix tsets

* introducing log crate; test for query router; comments

* typo

* fixes for banning

* added test for prepared stmt

Author: levkk

Cargo.lock

@@ -21,3 +21,5 @@ regex = "1"
 num_cpus = "1"
 once_cell = "1"
 statsd = "0.15"
+sqlparser = "0.14"
+log = "0.4"

Cargo.toml

@@ -81,3 +81,15 @@ database = "shard2"
 # replica: round-robin between replicas only without touching the primary,
 # primary: all queries go to the primary unless otherwise specified.
 default_role = "any"
+
+
+# Query parser. If enabled, we'll attempt to parse
+# every incoming query to determine if it's a read or a write.
+# If it's a read query, we'll direct it to a replica. Otherwise, if it's a write,
+# we'll direct it to the primary.
+query_parser_enabled = false
+
+# If the query parser is enabled and this setting is enabled, the primary will be part of the pool of databases used for
+# load balancing of read queries. Otherwise, the primary will only be used for write
+# queries. The primary can always be explicitely selected with our custom protocol.
+primary_reads_enabled = true

pgcat.toml

@@ -10,7 +10,6 @@ use tokio::net::{
 
 use std::collections::HashMap;
 
-use crate::config::Role;
 use crate::constants::*;
 use crate::errors::Error;
 use crate::messages::*;
@@ -47,10 +46,6 @@ pub struct Client {
     // to connect and cancel a query.
     client_server_map: ClientServerMap,
 
-    // Unless client specifies, route queries to the servers that have this role,
-    // e.g. primary or replicas or any.
-    default_server_role: Option<Role>,
-
     // Client parameters, e.g. user, client_encoding, etc.
     #[allow(dead_code)]
     parameters: HashMap<String, String>,
@@ -67,7 +62,6 @@ impl Client {
         mut stream: TcpStream,
         client_server_map: ClientServerMap,
         transaction_mode: bool,
-        default_server_role: Option<Role>,
         server_info: BytesMut,
         stats: Reporter,
     ) -> Result<Client, Error> {
@@ -126,7 +120,6 @@ impl Client {
                         process_id: process_id,
                         secret_key: secret_key,
                         client_server_map: client_server_map,
-                        default_server_role: default_server_role,
                         parameters: parameters,
                         stats: stats,
                     });
@@ -148,7 +141,6 @@ impl Client {
                         process_id: process_id,
                         secret_key: secret_key,
                         client_server_map: client_server_map,
-                        default_server_role: default_server_role,
                         parameters: HashMap::new(),
                         stats: stats,
                     });
@@ -162,7 +154,11 @@ impl Client {
     }
 
     /// Client loop. We handle all messages between the client and the database here.
-    pub async fn handle(&mut self, mut pool: ConnectionPool) -> Result<(), Error> {
+    pub async fn handle(
+        &mut self,
+        mut pool: ConnectionPool,
+        mut query_router: QueryRouter,
+    ) -> Result<(), Error> {
         // The client wants to cancel a query it has issued previously.
         if self.cancel_mode {
             let (process_id, secret_key, address, port) = {
@@ -191,8 +187,6 @@ impl Client {
             return Ok(Server::cancel(&address, &port, process_id, secret_key).await?);
         }
 
-        let mut query_router = QueryRouter::new(self.default_server_role, pool.shards());
-
         // Our custom protocol loop.
         // We expect the client to either start a transaction with regular queries
         // or issue commands for our sharding and server selection protocols.
@@ -222,6 +216,11 @@ impl Client {
                 continue;
             }
 
+            // Attempt to parse the query to determine where it should go
+            if query_router.query_parser_enabled() && query_router.role() == None {
+                query_router.infer_role(message.clone());
+            }
+
             // Grab a server from the pool: the client issued a regular query.
             let connection = match pool.get(query_router.shard(), query_router.role()).await {
                 Ok(conn) => conn,

src/client.rs

@@ -13,6 +13,24 @@ pub enum Role {
     Replica,
 }
 
+impl PartialEq<Option<Role>> for Role {
+    fn eq(&self, other: &Option<Role>) -> bool {
+        match other {
+            None => true,
+            Some(role) => *self == *role,
+        }
+    }
+}
+
+impl PartialEq<Role> for Option<Role> {
+    fn eq(&self, other: &Role) -> bool {
+        match *self {
+            None => true,
+            Some(role) => role == *other,
+        }
+    }
+}
+
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Debug)]
 pub struct Address {
     pub host: String,
@@ -47,6 +65,8 @@ pub struct Shard {
 #[derive(Deserialize, Debug, Clone)]
 pub struct QueryRouter {
     pub default_role: String,
+    pub query_parser_enabled: bool,
+    pub primary_reads_enabled: bool,
 }
 
 #[derive(Deserialize, Debug, Clone)]

src/config.rs

@@ -16,11 +16,13 @@
 extern crate async_trait;
 extern crate bb8;
 extern crate bytes;
+extern crate log;
 extern crate md5;
 extern crate num_cpus;
 extern crate once_cell;
 extern crate serde;
 extern crate serde_derive;
+extern crate sqlparser;
 extern crate statsd;
 extern crate tokio;
 extern crate toml;
@@ -47,6 +49,7 @@ mod stats;
 // secret keys to the backend's.
 use config::Role;
 use pool::{ClientServerMap, ConnectionPool};
+use query_router::QueryRouter;
 use stats::{Collector, Reporter};
 
 /// Main!
@@ -118,6 +121,8 @@ async fn main() {
             return;
         }
     };
+    let primary_reads_enabled = config.query_router.primary_reads_enabled;
+    let query_parser_enabled = config.query_router.query_parser_enabled;
 
     let server_info = match pool.validate().await {
         Ok(info) => info,
@@ -155,7 +160,6 @@ async fn main() {
                     socket,
                     client_server_map,
                     transaction_mode,
-                    default_server_role,
                     server_info,
                     reporter,
                 )
@@ -164,7 +168,14 @@ async fn main() {
                     Ok(mut client) => {
                         println!(">> Client {:?} authenticated successfully!", addr);
 
-                        match client.handle(pool).await {
+                        let query_router = QueryRouter::new(
+                            default_server_role,
+                            pool.shards(),
+                            primary_reads_enabled,
+                            query_parser_enabled,
+                        );
+
+                        match client.handle(pool, query_router).await {
                             Ok(()) => {
                                 let duration = chrono::offset::Utc::now().naive_utc() - start;