cufflink_pgcodec/
statement.rs

1//! What a guest's statement is, judged by its leading keywords: whatever
2//! runs a guest's SQL decides the same way how it may run.
3
4/// Statements Postgres refuses inside a transaction block, judged by their
5/// leading keywords after any comments: `CREATE`/`DROP INDEX CONCURRENTLY`,
6/// `REINDEX … CONCURRENTLY`, `REINDEX DATABASE`/`SYSTEM`, `VACUUM`,
7/// `CREATE`/`DROP DATABASE`, `CREATE`/`DROP TABLESPACE`, `ALTER SYSTEM`,
8/// `CLUSTER`, `DISCARD ALL`.
9pub fn requires_autocommit(sql: &str) -> bool {
10    let words = leading_keywords(sql, 8);
11    let word = |i: usize| words.get(i).map_or("", String::as_str);
12    let position = |keyword: &str| words.iter().position(|w| w == keyword);
13    match word(0) {
14        "VACUUM" | "CLUSTER" => true,
15        "DISCARD" => word(1) == "ALL",
16        "ALTER" => word(1) == "SYSTEM",
17        "CREATE" | "DROP" => {
18            matches!(word(1), "DATABASE" | "TABLESPACE")
19                || position("INDEX").is_some_and(|i| i <= 2 && word(i + 1) == "CONCURRENTLY")
20        }
21        "REINDEX" => words
22            .iter()
23            .position(|w| {
24                matches!(
25                    w.as_str(),
26                    "INDEX" | "TABLE" | "SCHEMA" | "DATABASE" | "SYSTEM"
27                )
28            })
29            .is_some_and(|kind| {
30                matches!(word(kind), "DATABASE" | "SYSTEM")
31                    || words[..=(kind + 1).min(words.len() - 1)]
32                        .iter()
33                        .any(|w| w == "CONCURRENTLY")
34            }),
35        _ => false,
36    }
37}
38
39/// Statements that change session state rather than the transaction's:
40/// `SET` other than `SET LOCAL`, `SET TRANSACTION` and `SET CONSTRAINTS`,
41/// and `RESET`.
42pub fn changes_session(sql: &str) -> bool {
43    let words = leading_keywords(sql, 2);
44    match words.first().map(String::as_str) {
45        Some("RESET") => true,
46        Some("SET") => !matches!(
47            words.get(1).map(String::as_str),
48            Some("LOCAL" | "TRANSACTION" | "CONSTRAINTS")
49        ),
50        _ => false,
51    }
52}
53
54/// Transaction control, which inside a host-managed transaction would end
55/// or split it behind the host's back.
56pub fn controls_transaction(sql: &str) -> bool {
57    let words = leading_keywords(sql, 2);
58    match words.first().map(String::as_str) {
59        Some(
60            "BEGIN" | "START" | "COMMIT" | "END" | "ROLLBACK" | "ABORT" | "SAVEPOINT" | "RELEASE",
61        ) => true,
62        Some("PREPARE") => words.get(1).is_some_and(|w| w == "TRANSACTION"),
63        _ => false,
64    }
65}
66
67/// A caller's `SELECT` wrapped to answer each row as JSON, `limit` rows at
68/// most. The SQL is wrapped, never concatenated onto another statement, and
69/// the newline after it keeps a query that ends in a `--` comment from
70/// swallowing the limit.
71pub fn rows_as_json(sql: &str, limit: i64) -> String {
72    format!("SELECT to_jsonb(t) AS row FROM (\n{sql}\n) t LIMIT {limit}")
73}
74
75/// The first `n` words of a statement, upper-cased, with leading `--` and
76/// `/* */` comments skipped and punctuation split off.
77fn leading_keywords(sql: &str, n: usize) -> Vec<String> {
78    let mut rest = sql;
79    loop {
80        rest = rest.trim_start();
81        if let Some(after) = rest.strip_prefix("--") {
82            rest = after.split_once('\n').map_or("", |(_, r)| r);
83        } else if let Some(after) = rest.strip_prefix("/*") {
84            rest = after.split_once("*/").map_or("", |(_, r)| r);
85        } else {
86            break;
87        }
88    }
89    rest.split(|c: char| c.is_whitespace() || matches!(c, '(' | ')' | ',' | ';'))
90        .filter(|t| !t.is_empty())
91        .take(n)
92        .map(|t| t.to_ascii_uppercase())
93        .collect()
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn a_trailing_comment_cannot_swallow_the_limit() {
102        assert_eq!(
103            rows_as_json("SELECT 1 -- all of them", 11),
104            "SELECT to_jsonb(t) AS row FROM (\nSELECT 1 -- all of them\n) t LIMIT 11"
105        );
106    }
107
108    #[test]
109    fn session_settings_are_told_from_transaction_ones() {
110        for sql in [
111            "SET search_path TO other",
112            "set app.x = '1'",
113            "SET SESSION statement_timeout = 0",
114            "SET ROLE admin",
115            "SET SESSION AUTHORIZATION admin",
116            "SET SESSION CHARACTERISTICS AS TRANSACTION READ ONLY",
117            "/* x */ SET TIME ZONE 'UTC'",
118            "RESET ALL",
119            "reset role",
120        ] {
121            assert!(changes_session(sql), "{sql}");
122        }
123        for sql in [
124            "SET LOCAL app.x = '1'",
125            "set local role reporter",
126            "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE",
127            "SET CONSTRAINTS ALL DEFERRED",
128            "SELECT set_config('app.x', '1', true)",
129            "UPDATE t SET a = 1",
130            "",
131        ] {
132            assert!(!changes_session(sql), "{sql}");
133        }
134    }
135
136    #[test]
137    fn transaction_control_is_recognised_by_leading_keywords() {
138        for sql in [
139            "BEGIN",
140            "begin isolation level serializable",
141            "START TRANSACTION",
142            "COMMIT",
143            "commit and chain",
144            "END",
145            "ROLLBACK",
146            "ROLLBACK TO SAVEPOINT sp_1",
147            "ABORT",
148            "SAVEPOINT x",
149            "RELEASE SAVEPOINT x",
150            "PREPARE TRANSACTION 'x'",
151            "COMMIT PREPARED 'x'",
152            "-- done\nCOMMIT;",
153        ] {
154            assert!(controls_transaction(sql), "{sql}");
155        }
156        for sql in [
157            "SELECT 1",
158            "PREPARE q AS SELECT 1",
159            "UPDATE t SET committed = true",
160            "WITH x AS (SELECT 1) SELECT * FROM x",
161            "",
162        ] {
163            assert!(!controls_transaction(sql), "{sql}");
164        }
165    }
166
167    #[test]
168    fn autocommit_statements_are_recognised_by_leading_keywords() {
169        for sql in [
170            "CREATE INDEX CONCURRENTLY idx ON t (c)",
171            "create unique index concurrently if not exists idx on t (c)",
172            "DROP INDEX CONCURRENTLY IF EXISTS idx",
173            "REINDEX INDEX CONCURRENTLY idx",
174            "REINDEX (VERBOSE) TABLE CONCURRENTLY t",
175            "REINDEX (CONCURRENTLY) TABLE t",
176            "REINDEX DATABASE mydb",
177            "reindex system",
178            "CREATE TABLESPACE fast LOCATION '/mnt/fast'",
179            "DROP TABLESPACE fast",
180            "  -- rebuild\n  VACUUM ANALYZE t",
181            "/* nightly */ VACUUM",
182            "CREATE DATABASE other",
183            "DROP DATABASE other",
184            "ALTER SYSTEM SET work_mem = '64MB'",
185            "CLUSTER t USING idx",
186            "DISCARD ALL",
187        ] {
188            assert!(requires_autocommit(sql), "{sql}");
189        }
190        for sql in [
191            "SELECT 1",
192            "INSERT INTO t VALUES (1)",
193            "CREATE INDEX idx ON t (c)",
194            "CREATE INDEX i ON t (concurrently)",
195            "REINDEX TABLE t",
196            "CREATE TABLE concurrently_log (id int)",
197            "ALTER TABLE t ADD COLUMN c int",
198            "-- CREATE INDEX CONCURRENTLY in a comment only\nSELECT 1",
199            "DISCARD TEMP",
200            "",
201        ] {
202            assert!(!requires_autocommit(sql), "{sql}");
203        }
204    }
205}
cufflink_pgcodec/statement.rs

cufflink_pgcodec/
statement.rs
