Plugins! (#420)

* Some queries

* Plugins!!

* cleanup

* actual names

* the actual plugins

* comment

* fix tests

* Tests

* unused errors

* Increase reaper rate to actually enforce settings

* ok

Author: levkk

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "1.0.1"
+version = "1.0.2-alpha1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -14,12 +14,12 @@ rand = "0.8"
 chrono = "0.4"
 sha-1 = "0.10"
 toml = "0.7"
-serde = "1"
+serde = { version = "1", features = ["derive"] }
 serde_derive = "1"
 regex = "1"
 num_cpus = "1"
 once_cell = "1"
-sqlparser = "0.33.0"
+sqlparser = {version = "0.33", features = ["visitor"] }
 log = "0.4"
 arc-swap = "1"
 env_logger = "0.10"
@@ -44,6 +44,7 @@ webpki-roots = "0.23"
 rustls = { version = "0.21", features = ["dangerous_configuration"] }
 trust-dns-resolver = "0.22.0"
 tokio-test = "0.4.2"
+serde_json = "1"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"

Cargo.toml

@@ -77,6 +77,9 @@ admin_username = "admin_user"
 # Password to access the virtual administrative database
 admin_password = "admin_pass"
 
+# Plugins!!
+# query_router_plugins = ["pg_table_access", "intercept"]
+
 # pool configs are structured as pool.<pool_name>
 # the pool_name is what clients use as database name when connecting.
 # For a pool named `sharded_db`, clients access that pool using connection string like

pgcat.toml

@@ -12,9 +12,9 @@ use tokio::time::Instant;
 use crate::config::{get_config, reload_config, VERSION};
 use crate::errors::Error;
 use crate::messages::*;
+use crate::pool::ClientServerMap;
 use crate::pool::{get_all_pools, get_pool};
 use crate::stats::{get_client_stats, get_pool_stats, get_server_stats, ClientState, ServerState};
-use crate::ClientServerMap;
 
 pub fn generate_server_info_for_admin() -> BytesMut {
     let mut server_info = BytesMut::new();

src/admin.rs

@@ -16,6 +16,7 @@ use crate::auth_passthrough::refetch_auth_hash;
 use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
 use crate::constants::*;
 use crate::messages::*;
+use crate::plugins::PluginOutput;
 use crate::pool::{get_pool, ClientServerMap, ConnectionPool};
 use crate::query_router::{Command, QueryRouter};
 use crate::server::Server;
@@ -765,6 +766,9 @@ where
 
         self.stats.register(self.stats.clone());
 
+        // Result returned by one of the plugins.
+        let mut plugin_output = None;
+
         // 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 protocol.
@@ -815,15 +819,39 @@ where
 
                 'Q' => {
                     if query_router.query_parser_enabled() {
-                        query_router.infer(&message);
+                        if let Ok(ast) = QueryRouter::parse(&message) {
+                            let plugin_result = query_router.execute_plugins(&ast).await;
+
+                            match plugin_result {
+                                Ok(PluginOutput::Deny(error)) => {
+                                    error_response(&mut self.write, &error).await?;
+                                    continue;
+                                }
+
+                                Ok(PluginOutput::Intercept(result)) => {
+                                    write_all(&mut self.write, result).await?;
+                                    continue;
+                                }
+
+                                _ => (),
+                            };
+
+                            let _ = query_router.infer(&ast);
+                        }
                     }
                 }
 
                 'P' => {
                     self.buffer.put(&message[..]);
 
                     if query_router.query_parser_enabled() {
-                        query_router.infer(&message);
+                        if let Ok(ast) = QueryRouter::parse(&message) {
+                            if let Ok(output) = query_router.execute_plugins(&ast).await {
+                                plugin_output = Some(output);
+                            }
+
+                            let _ = query_router.infer(&ast);
+                        }
                     }
 
                     continue;
@@ -857,6 +885,18 @@ where
                 continue;
             }
 
+            // Check on plugin results.
+            match plugin_output {
+                Some(PluginOutput::Deny(error)) => {
+                    self.buffer.clear();
+                    error_response(&mut self.write, &error).await?;
+                    plugin_output = None;
+                    continue;
+                }
+
+                _ => (),
+            };
+
             // Get a pool instance referenced by the most up-to-date
             // pointer. This ensures we always read the latest config
             // when starting a query.
@@ -1085,6 +1125,27 @@ where
                 match code {
                     // Query
                     'Q' => {
+                        if query_router.query_parser_enabled() {
+                            if let Ok(ast) = QueryRouter::parse(&message) {
+                                let plugin_result = query_router.execute_plugins(&ast).await;
+
+                                match plugin_result {
+                                    Ok(PluginOutput::Deny(error)) => {
+                                        error_response(&mut self.write, &error).await?;
+                                        continue;
+                                    }
+
+                                    Ok(PluginOutput::Intercept(result)) => {
+                                        write_all(&mut self.write, result).await?;
+                                        continue;
+                                    }
+
+                                    _ => (),
+                                };
+
+                                let _ = query_router.infer(&ast);
+                            }
+                        }
                         debug!("Sending query to server");
 
                         self.send_and_receive_loop(
@@ -1124,6 +1185,14 @@ where
                     // Parse
                     // The query with placeholders is here, e.g. `SELECT * FROM users WHERE email = $1 AND active = $2`.
                     'P' => {
+                        if query_router.query_parser_enabled() {
+                            if let Ok(ast) = QueryRouter::parse(&message) {
+                                if let Ok(output) = query_router.execute_plugins(&ast).await {
+                                    plugin_output = Some(output);
+                                }
+                            }
+                        }
+
                         self.buffer.put(&message[..]);
                     }
 
@@ -1155,6 +1224,24 @@ where
                     'S' => {
                         debug!("Sending query to server");
 
+                        match plugin_output {
+                            Some(PluginOutput::Deny(error)) => {
+                                error_response(&mut self.write, &error).await?;
+                                plugin_output = None;
+                                self.buffer.clear();
+                                continue;
+                            }
+
+                            Some(PluginOutput::Intercept(result)) => {
+                                write_all(&mut self.write, result).await?;
+                                plugin_output = None;
+                                self.buffer.clear();
+                                continue;
+                            }
+
+                            _ => (),
+                        };
+
                         self.buffer.put(&message[..]);
 
                         let first_message_code = (*self.buffer.get(0).unwrap_or(&0)) as char;

src/client.rs

@@ -298,9 +298,12 @@ pub struct General {
     pub admin_username: String,
     pub admin_password: String,
 
+    // Support for auth query
     pub auth_query: Option<String>,
     pub auth_query_user: Option<String>,
     pub auth_query_password: Option<String>,
+
+    pub query_router_plugins: Option<Vec<String>>,
 }
 
 impl General {
@@ -401,6 +404,7 @@ impl Default for General {
             auth_query_user: None,
             auth_query_password: None,
             server_lifetime: 1000 * 3600 * 24, // 24 hours,
+            query_router_plugins: None,
         }
     }
 }

src/config.rs

@@ -1,7 +1,7 @@
 //! Errors.
 
 /// Various errors.
-#[derive(Debug, PartialEq)]
+#[derive(Debug, PartialEq, Clone)]
 pub enum Error {
     SocketError(String),
     ClientSocketError(String, ClientIdentifier),
@@ -24,6 +24,8 @@ pub enum Error {
     ParseBytesError(String),
     AuthError(String),
     AuthPassthroughError(String),
+    UnsupportedStatement,
+    QueryRouterParserError(String),
 }
 
 #[derive(Clone, PartialEq, Debug)]

src/errors.rs

@@ -1,12 +1,17 @@
+pub mod admin;
 pub mod auth_passthrough;
+pub mod client;
 pub mod config;
 pub mod constants;
 pub mod dns_cache;
 pub mod errors;
 pub mod messages;
 pub mod mirrors;
 pub mod multi_logger;
+pub mod plugins;
 pub mod pool;
+pub mod prometheus;
+pub mod query_router;
 pub mod scram;
 pub mod server;
 pub mod sharding;