pg_pool_safe_query/lib.rs
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97
#[derive(Debug)]
pub struct TransactionPoolCompatibility {
is_compatible: bool,
reason: String,
statement_types: Vec<String>
}
impl TransactionPoolCompatibility {
pub fn new(is_compatible: bool, reason: String, statement_types: Vec<String>) -> TransactionPoolCompatibility {
TransactionPoolCompatibility {
is_compatible,
reason,
statement_types
}
}
// static method on the struct
pub fn parse(subject: String) -> TransactionPoolCompatibility {
let result = pg_query::parse(&subject).unwrap();
for (node, depth, parent) in result.protobuf.nodes() {
// println!("{} node:{:?} parent:{:?}", " ".repeat((depth as usize) * 2), node, parent);
match node {
pg_query::NodeRef::VariableSetStmt(stmt) => println!("got em! {:?} {:?}", stmt, stmt.is_local),
_ => println!("nope"),
}
}
TransactionPoolCompatibility {
is_compatible: true,
reason: "it's compatible".to_string(),
statement_types: vec!["SELECT".to_string(), "INSERT".to_string()]
}
}
}
pub fn is_pgbouncer_compatible(subject: String) -> String {
let result = pg_query::parse(&subject.to_string());
assert!(result.is_ok());
let result = result.unwrap();
println!("{:?}", result.statement_types());
result.statement_types()[0].to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn it_is_pgbouncer_compatible() {
// println!(
// "{:?}",
// pg_query::parse("SELECT * FROM contacts; SET statement_timeout TO '1s';").unwrap().statement_types()
// );
// println!(
// "{:?}",
// pg_query::parse("SELECT * FROM contacts; SET statement_timeout TO '1s';").unwrap().protobuf.nodes()
// );
let result = TransactionPoolCompatibility::parse(
"SELECT * FROM contacts; SET statement_timeout TO '1s';".to_string()
);
println!("{:?}", result);
// for (node, depth, parent) in result.protobuf.nodes() {
// println!("{} node:{:?} parent:{:?}", " ".repeat((depth as usize) * 2), node, parent);
// match node {
// pg_query::NodeRef::VariableSetStmt(stmt) => println!("got em! {:?} {:?}", stmt, stmt.is_local),
// _ => println!("nope"),
// }
// }
// println!(
// "{:?}",
// pg_query::parse("SELECT * FROM contacts; SET statement_timeout TO '1s';").unwrap().protobuf.stmts
// );
// pg_query::parse("SELECT * FROM contacts; SET statement_timeout TO '1s';").unwrap();
// println!("{:?}", pg_query::split_with_parser("SELECT * FROM contacts; SET statement_timeout TO '1s';"));
// let result = is_pgbouncer_compatible("SELECT * FROM contacts".to_string());
// assert_eq!(result, "SELECT");
// let mut i = 0;
// while i < 100000 {
// i += 1;
// pg_query::parse("SET LOCAL lock_timeout TO '10s'").unwrap();
// }
assert_eq!(is_pgbouncer_compatible("SET LOCAL lock_timeout TO '10s';".to_string()), "VariableSetStmt");
assert_eq!(is_pgbouncer_compatible("LISTEN channel_name;".to_string()), "ListenStmt");
assert_eq!(is_pgbouncer_compatible("UNLISTEN channel_name;".to_string()), "UnlistenStmt");
assert_eq!(is_pgbouncer_compatible("BEGIN;".to_string()), "TransactionStmt");
assert_eq!(is_pgbouncer_compatible("COMMIT;".to_string()), "TransactionStmt");
assert_eq!(is_pgbouncer_compatible("ROLLBACK;".to_string()), "TransactionStmt");
assert_eq!(is_pgbouncer_compatible("PREPARE prepared_statement AS SELECT * FROM my_table;".to_string()), "PrepareStmt");
assert_eq!(is_pgbouncer_compatible("EXECUTE prepared_statement;".to_string()), "ExecuteStmt");
assert_eq!(is_pgbouncer_compatible("DEALLOCATE prepared_statement;".to_string()), "DeallocateStmt");
}
}