docker/pgcat
1
0
Fork 0
mirror of https://github.com/postgresml/pgcat.git synced 2026-03-23 01:16:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

more plugins

Browse Source
This commit is contained in:
Lev Kokotov
2023-05-03 15:29:16 -07:00
parent d5e329fec5
commit 374a6b138b
8 changed files with 249 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
pgcat.toml
View File

@@ -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,33 @@ 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 = false
[plugins.intercept.queries.0]
query = "select current_database() as a, current_schemas(false) as b"
schema = [
["a", "text"],
["b", "text"],
]
result = [
["${DATABASE}", "{public}"],
]
# 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]

71
src/config.rs
View File

@@ -483,6 +483,7 @@ pub struct Pool {
pub auth_query: Option<String>,
pub auth_query_user: Option<String>,
pub auth_query_password: Option<String>,
pub query_router_plugins: Option<Vec<String>>,
pub shards: BTreeMap<String, Shard>,
pub users: BTreeMap<String, User>,
@@ -604,6 +605,7 @@ impl Default for Pool {
auth_query_user: None,
auth_query_password: None,
server_lifetime: None,
query_router_plugins: None,
}
}
}
@@ -682,6 +684,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 +751,7 @@ pub struct Config {
pub path: String,
pub general: General,
pub plugins: Option<Plugins>,
pub pools: HashMap<String, Pool>,
}
@@ -729,6 +781,14 @@ impl Config {
}
}
}
pub fn fill_up_query_router_plugins(&mut self) {
for (_name, pool) in self.pools.iter_mut() {
if pool.query_router_plugins.is_none() {
pool.query_router_plugins = self.general.query_router_plugins.clone();
}
}
}
}
impl Default for Config {
@@ -737,6 +797,7 @@ impl Default for Config {
path: Self::default_path(),
general: General::default(),
pools: HashMap::default(),
plugins: None,
}
}
}
@@ -1116,6 +1177,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
};
config.fill_up_auth_query_config();
config.fill_up_query_router_plugins();
config.validate()?;
config.path = path.to_string();
@@ -1128,6 +1190,7 @@ 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) => {
@@ -1135,18 +1198,18 @@ pub async fn reload_config(client_server_map: ClientServerMap) -> Result<bool, E
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)
}

49
src/plugins/intercept.rs
View File

@@ -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

9
src/plugins/mod.rs
View File

@@ -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);
}

49
src/plugins/query_logger.rs Normal file
View File

@@ -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)
}
}

31
src/plugins/table_access.rs
View File

@@ -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

29
src/pool.rs
View File

@@ -453,7 +453,7 @@ impl ConnectionPool {
auth_query: pool_config.auth_query.clone(),
auth_query_user: pool_config.auth_query_user.clone(),
auth_query_password: pool_config.auth_query_password.clone(),
plugins: config.general.query_router_plugins.clone(),
plugins: pool_config.query_router_plugins.clone(),
},
validated: Arc::new(AtomicBool::new(false)),
paused: Arc::new(AtomicBool::new(false)),
@@ -473,10 +473,29 @@ impl ConnectionPool {
}
}
// Initialize plugins here if required.
if let Some(plugins) = config.general.query_router_plugins {
if plugins.contains(&String::from("intercept")) {
crate::plugins::intercept::configure(&new_pools);
if let Some(ref plugins) = config.plugins {
if let Some(ref intercept) = plugins.intercept {
if intercept.enabled {
crate::plugins::intercept::setup(intercept, &new_pools);
} else {
crate::plugins::intercept::disable();
}
}
if let Some(ref table_access) = plugins.table_access {
if table_access.enabled {
crate::plugins::table_access::setup(table_access);
} else {
crate::plugins::table_access::disable();
}
}
if let Some(ref query_logger) = plugins.query_logger {
if query_logger.enabled {
crate::plugins::query_logger::setup();
} else {
crate::plugins::query_logger::disable();
}
}
}

35
src/query_router.rs
View File

@@ -15,7 +15,10 @@ use sqlparser::parser::Parser;
use crate::config::Role;
use crate::errors::Error;
use crate::messages::BytesMutReader;
use crate::plugins::{Intercept, Plugin, PluginOutput, TableAccess};
use crate::plugins::{
intercept, query_logger, table_access, Intercept, Plugin, PluginOutput, QueryLogger,
TableAccess,
};
use crate::pool::PoolSettings;
use crate::sharding::Sharder;
@@ -790,24 +793,26 @@ impl QueryRouter {
/// Add your plugins here and execute them.
pub async fn execute_plugins(&self, ast: &Vec<Statement>) -> Result<PluginOutput, Error> {
if let Some(plugins) = &self.pool_settings.plugins {
if plugins.contains(&String::from("intercept")) {
let mut intercept = Intercept {};
let result = intercept.run(&self, ast).await;
if query_logger::enabled() {
let mut query_logger = QueryLogger {};
let _ = query_logger.run(&self, ast).await;
}
if let Ok(PluginOutput::Intercept(output)) = result {
return Ok(PluginOutput::Intercept(output));
}
if intercept::enabled() {
let mut intercept = Intercept {};
let result = intercept.run(&self, ast).await;
if let Ok(PluginOutput::Intercept(output)) = result {
return Ok(PluginOutput::Intercept(output));
}
}
if plugins.contains(&String::from("pg_table_access")) {
let mut table_access = TableAccess {
forbidden_tables: vec![String::from("pg_database"), String::from("pg_roles")],
};
if table_access::enabled() {
let mut table_access = TableAccess {};
let result = table_access.run(&self, ast).await;
if let Ok(PluginOutput::Deny(error)) = table_access.run(&self, ast).await {
return Ok(PluginOutput::Deny(error));
}
if let Ok(PluginOutput::Deny(error)) = result {
return Ok(PluginOutput::Deny(error));
}
}