Prewarmer (#435)

* Prewarmer

* hmm

* Tests

* default

* fix test

* Correct configuration

* Added minimal config example

* remove connect_timeout

src/query_router.rs

@@ -15,10 +15,7 @@ use sqlparser::parser::Parser;
 use crate::config::Role;
 use crate::errors::Error;
 use crate::messages::BytesMutReader;
-use crate::plugins::{
-    intercept, query_logger, table_access, Intercept, Plugin, PluginOutput, QueryLogger,
-    TableAccess,
-};
+use crate::plugins::{Intercept, Plugin, PluginOutput, QueryLogger, TableAccess};
 use crate::pool::PoolSettings;
 use crate::sharding::Sharder;
 
@@ -793,13 +790,27 @@ impl QueryRouter {
 
     /// Add your plugins here and execute them.
     pub async fn execute_plugins(&self, ast: &Vec<Statement>) -> Result<PluginOutput, Error> {
-        if query_logger::enabled() {
-            let mut query_logger = QueryLogger {};
+        let plugins = match self.pool_settings.plugins {
+            Some(ref plugins) => plugins,
+            None => return Ok(PluginOutput::Allow),
+        };
+
+        if let Some(ref query_logger) = plugins.query_logger {
+            let mut query_logger = QueryLogger {
+                enabled: query_logger.enabled,
+                user: &self.pool_settings.user.username,
+                db: &self.pool_settings.db,
+            };
+
             let _ = query_logger.run(&self, ast).await;
         }
 
-        if intercept::enabled() {
-            let mut intercept = Intercept {};
+        if let Some(ref intercept) = plugins.intercept {
+            let mut intercept = Intercept {
+                enabled: intercept.enabled,
+                config: &intercept,
+            };
+
             let result = intercept.run(&self, ast).await;
 
             if let Ok(PluginOutput::Intercept(output)) = result {
@@ -807,8 +818,12 @@ impl QueryRouter {
             }
         }
 
-        if table_access::enabled() {
-            let mut table_access = TableAccess {};
+        if let Some(ref table_access) = plugins.table_access {
+            let mut table_access = TableAccess {
+                enabled: table_access.enabled,
+                tables: &table_access.tables,
+            };
+
             let result = table_access.run(&self, ast).await;
 
             if let Ok(PluginOutput::Deny(error)) = result {
@@ -1161,6 +1176,7 @@ mod test {
             auth_query_password: None,
             auth_query_user: None,
             db: "test".to_string(),
+            plugins: None,
         };
         let mut qr = QueryRouter::new();
         assert_eq!(qr.active_role, None);
@@ -1235,7 +1251,9 @@ mod test {
             auth_query_password: None,
             auth_query_user: None,
             db: "test".to_string(),
+            plugins: None,
         };
+
         let mut qr = QueryRouter::new();
         qr.update_pool_settings(pool_settings.clone());
 
@@ -1379,17 +1397,25 @@ mod test {
 
     #[tokio::test]
     async fn test_table_access_plugin() {
-        use crate::config::TableAccess;
-        let ta = TableAccess {
+        use crate::config::{Plugins, TableAccess};
+        let table_access = TableAccess {
             enabled: true,
             tables: vec![String::from("pg_database")],
         };
-
-        crate::plugins::table_access::setup(&ta);
+        let plugins = Plugins {
+            table_access: Some(table_access),
+            intercept: None,
+            query_logger: None,
+            prewarmer: None,
+        };
 
         QueryRouter::setup();
+        let mut pool_settings = PoolSettings::default();
+        pool_settings.query_parser_enabled = true;
+        pool_settings.plugins = Some(plugins);
 
-        let qr = QueryRouter::new();
+        let mut qr = QueryRouter::new();
+        qr.update_pool_settings(pool_settings);
 
         let query = simple_query("SELECT * FROM pg_database");
         let ast = QueryRouter::parse(&query).unwrap();
@@ -1403,4 +1429,17 @@ mod test {
             ))
         );
     }
+
+    #[tokio::test]
+    async fn test_plugins_disabled_by_defaault() {
+        QueryRouter::setup();
+        let qr = QueryRouter::new();
+
+        let query = simple_query("SELECT * FROM pg_database");
+        let ast = QueryRouter::parse(&query).unwrap();
+
+        let res = qr.execute_plugins(&ast).await;
+
+        assert_eq!(res, Ok(PluginOutput::Allow));
+    }
 }