pg_pool_safe_query/
lib.rs

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

pg_pool_safe_query/
lib.rs
