Actually plugins (#421)

* more plugins

* clean up

* fix tests

* fix flakey test

src/plugins/intercept.rs

@@ -11,10 +11,11 @@ use serde_json::{json, Value};
 use sqlparser::ast::Statement;
 use std::collections::HashMap;
 
-use log::debug;
+use log::{debug, info};
 use std::sync::Arc;
 
 use crate::{
+    config::Intercept as InterceptConfig,
     errors::Error,
     messages::{command_complete, data_row_nullable, row_description, DataType},
     plugins::{Plugin, PluginOutput},
@@ -22,19 +23,29 @@ use crate::{
     query_router::QueryRouter,
 };
 
-pub static CONFIG: Lazy<ArcSwap<HashMap<PoolIdentifier, Value>>> =
+pub static CONFIG: Lazy<ArcSwap<HashMap<PoolIdentifier, InterceptConfig>>> =
     Lazy::new(|| ArcSwap::from_pointee(HashMap::new()));
 
-/// Configure the intercept plugin.
-pub fn configure(pools: &PoolMap) {
+/// Check if the interceptor plugin has been enabled.
+pub fn enabled() -> bool {
+    !CONFIG.load().is_empty()
+}
+
+pub fn setup(intercept_config: &InterceptConfig, pools: &PoolMap) {
     let mut config = HashMap::new();
     for (identifier, _) in pools.iter() {
-        // TODO: make this configurable from a text config.
-        let value = fool_datagrip(&identifier.db, &identifier.user);
-        config.insert(identifier.clone(), value);
+        let mut intercept_config = intercept_config.clone();
+        intercept_config.substitute(&identifier.db, &identifier.user);
+        config.insert(identifier.clone(), intercept_config);
     }
 
     CONFIG.store(Arc::new(config));
+
+    info!("Intercepting {} queries", intercept_config.queries.len());
+}
+
+pub fn disable() {
+    CONFIG.store(Arc::new(HashMap::new()));
 }
 
 // TODO: use these structs for deserialization
@@ -78,19 +89,19 @@ impl Plugin for Intercept {
             // Normalization
             let q = q.to_string().to_ascii_lowercase();
 
-            for target in query_map.as_array().unwrap().iter() {
-                if target["query"].as_str().unwrap() == q {
-                    debug!("Query matched: {}", q);
+            for (_, target) in query_map.queries.iter() {
+                if target.query.as_str() == q {
+                    debug!("Intercepting query: {}", q);
 
-                    let rd = target["schema"]
-                        .as_array()
-                        .unwrap()
+                    let rd = target
+                        .schema
                         .iter()
                         .map(|row| {
-                            let row = row.as_object().unwrap();
+                            let name = &row[0];
+                            let data_type = &row[1];
                             (
-                                row["name"].as_str().unwrap(),
-                                match row["data_type"].as_str().unwrap() {
+                                name.as_str(),
+                                match data_type.as_str() {
                                     "text" => DataType::Text,
                                     "anyarray" => DataType::AnyArray,
                                     "oid" => DataType::Oid,
@@ -104,13 +115,11 @@ impl Plugin for Intercept {
 
                     result.put(row_description(&rd));
 
-                    target["result"].as_array().unwrap().iter().for_each(|row| {
+                    target.result.iter().for_each(|row| {
                         let row = row
-                            .as_array()
-                            .unwrap()
                             .iter()
                             .map(|s| {
-                                let s = s.as_str().unwrap().to_string();
+                                let s = s.as_str().to_string();
 
                                 if s == "" {
                                     None
@@ -141,6 +150,7 @@ impl Plugin for Intercept {
 
 /// Make IntelliJ SQL plugin believe it's talking to an actual database
 /// instead of PgCat.
+#[allow(dead_code)]
 fn fool_datagrip(database: &str, user: &str) -> Value {
     json!([
         {

src/plugins/mod.rs

@@ -9,6 +9,7 @@
 //!
 
 pub mod intercept;
+pub mod query_logger;
 pub mod table_access;
 
 use crate::{errors::Error, query_router::QueryRouter};
@@ -17,6 +18,7 @@ use bytes::BytesMut;
 use sqlparser::ast::Statement;
 
 pub use intercept::Intercept;
+pub use query_logger::QueryLogger;
 pub use table_access::TableAccess;
 
 #[derive(Clone, Debug, PartialEq)]
@@ -29,12 +31,13 @@ pub enum PluginOutput {
 
 #[async_trait]
 pub trait Plugin {
-    // Custom output is allowed because we want to extend this system
-    // to rewriting queries some day. So an output of a plugin could be
-    // a rewritten AST.
+    // Run before the query is sent to the server.
     async fn run(
         &mut self,
         query_router: &QueryRouter,
         ast: &Vec<Statement>,
     ) -> Result<PluginOutput, Error>;
+
+    // TODO: run after the result is returned
+    // async fn callback(&mut self, query_router: &QueryRouter);
 }

src/plugins/query_logger.rs

@@ -0,0 +1,49 @@
+//! Log all queries to stdout (or somewhere else, why not).
+
+use crate::{
+    errors::Error,
+    plugins::{Plugin, PluginOutput},
+    query_router::QueryRouter,
+};
+use arc_swap::ArcSwap;
+use async_trait::async_trait;
+use log::info;
+use once_cell::sync::Lazy;
+use sqlparser::ast::Statement;
+use std::sync::Arc;
+
+static ENABLED: Lazy<ArcSwap<bool>> = Lazy::new(|| ArcSwap::from_pointee(false));
+
+pub struct QueryLogger;
+
+pub fn setup() {
+    ENABLED.store(Arc::new(true));
+
+    info!("Logging queries to stdout");
+}
+
+pub fn disable() {
+    ENABLED.store(Arc::new(false));
+}
+
+pub fn enabled() -> bool {
+    **ENABLED.load()
+}
+
+#[async_trait]
+impl Plugin for QueryLogger {
+    async fn run(
+        &mut self,
+        _query_router: &QueryRouter,
+        ast: &Vec<Statement>,
+    ) -> Result<PluginOutput, Error> {
+        let query = ast
+            .iter()
+            .map(|q| q.to_string())
+            .collect::<Vec<String>>()
+            .join("; ");
+        info!("{}", query);
+
+        Ok(PluginOutput::Allow)
+    }
+}

src/plugins/table_access.rs

@@ -5,17 +5,37 @@ use async_trait::async_trait;
 use sqlparser::ast::{visit_relations, Statement};
 
 use crate::{
+    config::TableAccess as TableAccessConfig,
     errors::Error,
     plugins::{Plugin, PluginOutput},
     query_router::QueryRouter,
 };
 
-use core::ops::ControlFlow;
+use log::{debug, info};
 
-pub struct TableAccess {
-    pub forbidden_tables: Vec<String>,
+use arc_swap::ArcSwap;
+use core::ops::ControlFlow;
+use once_cell::sync::Lazy;
+use std::sync::Arc;
+
+static CONFIG: Lazy<ArcSwap<Vec<String>>> = Lazy::new(|| ArcSwap::from_pointee(vec![]));
+
+pub fn setup(config: &TableAccessConfig) {
+    CONFIG.store(Arc::new(config.tables.clone()));
+
+    info!("Blocking access to {} tables", config.tables.len());
 }
 
+pub fn enabled() -> bool {
+    !CONFIG.load().is_empty()
+}
+
+pub fn disable() {
+    CONFIG.store(Arc::new(vec![]));
+}
+
+pub struct TableAccess;
+
 #[async_trait]
 impl Plugin for TableAccess {
     async fn run(
@@ -24,13 +44,14 @@ impl Plugin for TableAccess {
         ast: &Vec<Statement>,
     ) -> Result<PluginOutput, Error> {
         let mut found = None;
+        let forbidden_tables = CONFIG.load();
 
         visit_relations(ast, |relation| {
             let relation = relation.to_string();
             let parts = relation.split(".").collect::<Vec<&str>>();
             let table_name = parts.last().unwrap();
 
-            if self.forbidden_tables.contains(&table_name.to_string()) {
+            if forbidden_tables.contains(&table_name.to_string()) {
                 found = Some(table_name.to_string());
                 ControlFlow::<()>::Break(())
             } else {
@@ -39,6 +60,8 @@ impl Plugin for TableAccess {
         });
 
         if let Some(found) = found {
+            debug!("Blocking access to table \"{}\"", found);
+
             Ok(PluginOutput::Deny(format!(
                 "permission for table \"{}\" denied",
                 found