Actually plugins (#421)

* more plugins * clean up * fix tests * fix flakey test

Author: levkk

Cargo.toml

@@ -48,3 +48,4 @@ serde_json = "1"
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 jemallocator = "0.5.0"
+

pgcat.toml

@@ -77,9 +77,6 @@ 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
@@ -157,6 +154,45 @@ connect_timeout = 3000
 # Specifies how often (in seconds) cached ip addresses for servers are rechecked (see `dns_cache_enabled`).
 # dns_max_ttl = 30
 
+[plugins]
+
+[plugins.query_logger]
+enabled = false
+
+[plugins.table_access]
+enabled = false
+tables = [
+ "pg_user",
+ "pg_roles",
+ "pg_database",
+]
+
+[plugins.intercept]
+enabled = true
+
+[plugins.intercept.queries.0]
+
+query = "select current_database() as a, current_schemas(false) as b"
+schema = [
+ ["a", "text"],
+ ["b", "text"],
+]
+result = [
+ ["${DATABASE}", "{public}"],
+]
+
+[plugins.intercept.queries.1]
+
+query = "select current_database(), current_schema(), current_user"
+schema = [
+ ["current_database", "text"],
+ ["current_schema", "text"],
+ ["current_user", "text"],
+]
+result = [
+ ["${DATABASE}", "public", "${USER}"],
+]
+
 # User configs are structured as pool.<pool_name>.users.<user_index>
 # This section holds the credentials for users that may connect to this cluster
 [pools.sharded_db.users.0]

src/config.rs

@@ -302,8 +302,6 @@ pub struct General {
  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 {
@@ -404,7 +402,6 @@ impl Default for General {
  auth_query_user: None,
  auth_query_password: None,
  server_lifetime: 1000 * 3600 * 24, // 24 hours,
- query_router_plugins: None,
  }
  }
 }
@@ -682,6 +679,55 @@ impl Default for Shard {
  }
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct Plugins {
+ pub intercept: Option<Intercept>,
+ pub table_access: Option<TableAccess>,
+ pub query_logger: Option<QueryLogger>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct Intercept {
+ pub enabled: bool,
+ pub queries: BTreeMap<String, Query>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct TableAccess {
+ pub enabled: bool,
+ pub tables: Vec<String>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct QueryLogger {
+ pub enabled: bool,
+}
+
+impl Intercept {
+ pub fn substitute(&mut self, db: &str, user: &str) {
+ for (_, query) in self.queries.iter_mut() {
+ query.substitute(db, user);
+ }
+ }
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default)]
+pub struct Query {
+ pub query: String,
+ pub schema: Vec<Vec<String>>,
+ pub result: Vec<Vec<String>>,
+}
+
+impl Query {
+ pub fn substitute(&mut self, db: &str, user: &str) {
+ for col in self.result.iter_mut() {
+ for i in 0..col.len() {
+ col[i] = col[i].replace("${USER}", user).replace("${DATABASE}", db);
+ }
+ }
+ }
+}
+
 /// Configuration wrapper.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct Config {
@@ -700,6 +746,7 @@ pub struct Config {
  pub path: String,
 
  pub general: General,
+ pub plugins: Option<Plugins>,
  pub pools: HashMap<String, Pool>,
 }
 
@@ -737,6 +784,7 @@ impl Default for Config {
  path: Self::default_path(),
  general: General::default(),
  pools: HashMap::default(),
+ plugins: None,
  }
  }
 }
@@ -1128,25 +1176,26 @@ pub async fn parse(path: &str) -> Result<(), Error> {
 
 pub async fn reload_config(client_server_map: ClientServerMap) -> Result<bool, Error> {
  let old_config = get_config();
+
  match parse(&old_config.path).await {
  Ok(()) => (),
  Err(err) => {
  error!("Config reload error: {:?}", err);
  return Err(Error::BadConfig);
  }
  };
+
  let new_config = get_config();
+
  match CachedResolver::from_config().await {
  Ok(_) => (),
  Err(err) => error!("DNS cache reinitialization error: {:?}", err),
  };
 
- if old_config.pools != new_config.pools {
- info!("Pool configuration changed");
+ if old_config != new_config {
+ info!("Config changed, reloading");
  ConnectionPool::from_config(client_server_map).await?;
  Ok(true)
- } else if old_config != new_config {
- Ok(true)
  } else {
  Ok(false)
  }

src/plugins/intercept.rs

@@ -11,30 +11,41 @@ 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},
  pool::{PoolIdentifier, PoolMap},
  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)
+ }
+}