cufflink_pgcodec/
pooler.rs

1//! What sits between a pool and Postgres, and the guest statements a
2//! transaction pooler cannot carry.
3
4use crate::statement::{changes_session, requires_autocommit};
5use std::str::FromStr;
6
7/// What sits between a pool and Postgres.
8#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
9pub enum Pooler {
10    #[default]
11    Direct,
12    /// PgBouncer in `pool_mode=transaction`: a server connection is held
13    /// only for the length of a transaction, so no session state may be
14    /// relied on across statements, and startup parameters are refused.
15    PgBouncerTransaction,
16}
17
18impl FromStr for Pooler {
19    type Err = String;
20
21    fn from_str(s: &str) -> Result<Self, Self::Err> {
22        match s.trim().to_ascii_lowercase().as_str() {
23            "direct" => Ok(Pooler::Direct),
24            "pgbouncer_transaction" => Ok(Pooler::PgBouncerTransaction),
25            other => Err(format!(
26                "'{other}' is not one of direct, pgbouncer_transaction"
27            )),
28        }
29    }
30}
31
32impl Pooler {
33    /// Whether consecutive statements outside a transaction may land on
34    /// different server connections.
35    pub fn is_transaction_pooler(self) -> bool {
36        self == Pooler::PgBouncerTransaction
37    }
38}
39
40/// Why a transaction pooler refuses a guest statement.
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum PoolerRefusal {
43    /// It cannot run in a transaction, so its `SET` and `RESET` would land
44    /// on other server connections; such DDL belongs in deploy-time
45    /// migrations.
46    Autocommit,
47    /// It sets session state, which would outlive the transaction on a
48    /// server connection other clients share. `set_config(…, false)` and
49    /// `DO` blocks cannot be recognised this way.
50    SessionState,
51}
52
53impl PoolerRefusal {
54    pub fn message(self) -> &'static str {
55        match self {
56            PoolerRefusal::Autocommit => {
57                "a statement that cannot run inside a transaction is refused under \
58                 TENANT_DB_POOLER=pgbouncer_transaction; run it in a deploy-time migration"
59            }
60            PoolerRefusal::SessionState => {
61                "SET and RESET are refused under TENANT_DB_POOLER=pgbouncer_transaction; \
62                 use SET LOCAL inside a transaction"
63            }
64        }
65    }
66}
67
68/// Why `pooler` refuses the guest statement `sql`, if it does.
69pub fn pooler_refusal(pooler: Pooler, sql: &str) -> Option<PoolerRefusal> {
70    if !pooler.is_transaction_pooler() {
71        None
72    } else if requires_autocommit(sql) {
73        Some(PoolerRefusal::Autocommit)
74    } else if changes_session(sql) {
75        Some(PoolerRefusal::SessionState)
76    } else {
77        None
78    }
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    #[test]
86    fn a_pooler_is_named_in_any_case() {
87        assert_eq!(
88            "PgBouncer_Transaction".parse::<Pooler>(),
89            Ok(Pooler::PgBouncerTransaction)
90        );
91        assert_eq!(" direct ".parse::<Pooler>(), Ok(Pooler::Direct));
92        assert!("session".parse::<Pooler>().is_err());
93        assert!(Pooler::PgBouncerTransaction.is_transaction_pooler());
94        assert!(!Pooler::Direct.is_transaction_pooler());
95        assert_eq!(Pooler::default(), Pooler::Direct);
96    }
97
98    #[test]
99    fn a_transaction_pooler_refuses_autocommit_and_session_statements_only() {
100        let pooled = |sql| pooler_refusal(Pooler::PgBouncerTransaction, sql);
101        let concurrently = "CREATE INDEX CONCURRENTLY i ON t (c)";
102        assert_eq!(pooled(concurrently), Some(PoolerRefusal::Autocommit));
103        assert_eq!(pooled("SET ROLE other"), Some(PoolerRefusal::SessionState));
104        assert_eq!(pooled("SELECT 1"), None);
105        assert_eq!(pooled("SET LOCAL app.x = '1'"), None);
106        for sql in [concurrently, "SET ROLE other", "RESET ALL"] {
107            assert_eq!(pooler_refusal(Pooler::Direct, sql), None, "{sql}");
108        }
109    }
110}
cufflink_pgcodec/pooler.rs

cufflink_pgcodec/
pooler.rs
