reddb_wire/
sql_routing.rs

1//! Conservative client classification. Unknown SQL stays on the writer.
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum SqlRoute {
5    Replica,
6    Primary,
7    Begin,
8    End,
9    Session,
10}
11
12pub fn sql_route(sql: &str) -> SqlRoute {
13    let Some(mut tokens) = sql_tokens(sql) else {
14        return SqlRoute::Primary;
15    };
16    if tokens.last().is_some_and(|token| token == ";") {
17        tokens.pop();
18    }
19    let Some(first) = tokens.first().map(String::as_str) else {
20        return SqlRoute::Primary;
21    };
22    if first == "BEGIN"
23        || first == "START" && tokens.get(1).is_some_and(|token| token == "TRANSACTION")
24    {
25        return SqlRoute::Begin;
26    }
27    if matches!(first, "SET" | "USE" | "PREPARE" | "DECLARE" | "LISTEN") {
28        return SqlRoute::Session;
29    }
30    if tokens.iter().any(|token| token == ";") {
31        return SqlRoute::Session;
32    }
33    if matches!(first, "COMMIT" | "ROLLBACK") && tokens.iter().any(|token| token == "CHAIN") {
34        return SqlRoute::Begin;
35    }
36    if first == "COMMIT" || first == "ROLLBACK" && !tokens.iter().any(|token| token == "TO") {
37        return SqlRoute::End;
38    }
39    if first != "SELECT" {
40        return SqlRoute::Primary;
41    }
42    for (index, token) in tokens.iter().enumerate() {
43        if matches!(
44            token.as_str(),
45            "INSERT"
46                | "UPDATE"
47                | "DELETE"
48                | "CREATE"
49                | "ALTER"
50                | "DROP"
51                | "INTO"
52                | "FOR"
53                | "LOCK"
54                | "CALL"
55                | "EXECUTE"
56                | "SET"
57                | "QUEUE"
58                | "COPY"
59                | "["
60                | "]"
61        ) {
62            return SqlRoute::Primary;
63        }
64        if tokens.get(index + 1).is_some_and(|next| next == "(")
65            && token
66                .as_bytes()
67                .first()
68                .is_some_and(|byte| byte.is_ascii_alphabetic() || *byte == b'_' || *byte == b'#')
69            && !matches!(
70                token.as_str(),
71                "COUNT"
72                    | "SUM"
73                    | "MIN"
74                    | "MAX"
75                    | "AVG"
76                    | "LOWER"
77                    | "UPPER"
78                    | "LENGTH"
79                    | "COALESCE"
80                    | "ABS"
81                    | "ROUND"
82                    | "CAST"
83                    | "IN"
84                    | "EXISTS"
85                    | "SELECT"
86                    | "WHERE"
87                    | "AND"
88                    | "OR"
89                    | "NOT"
90                    | "FROM"
91                    | "#value"
92            )
93        {
94            return SqlRoute::Primary;
95        }
96    }
97    SqlRoute::Replica
98}
99
100fn sql_tokens(sql: &str) -> Option<Vec<String>> {
101    if sql.len() > 1_048_576 {
102        return None;
103    }
104    let bytes = sql.as_bytes();
105    let mut index = 0;
106    let mut tokens = Vec::new();
107    while index < bytes.len() {
108        let byte = bytes[index];
109        if byte.is_ascii_whitespace() {
110            index += 1;
111            continue;
112        }
113        if bytes[index..].starts_with(b"--") {
114            index += 2;
115            while index < bytes.len() && bytes[index] != b'\n' {
116                index += 1;
117            }
118            continue;
119        }
120        if bytes[index..].starts_with(b"/*") {
121            let mut depth = 1;
122            index += 2;
123            while index < bytes.len() && depth > 0 {
124                if bytes[index..].starts_with(b"/*") {
125                    depth += 1;
126                    index += 2;
127                } else if bytes[index..].starts_with(b"*/") {
128                    depth -= 1;
129                    index += 2;
130                } else {
131                    index += 1;
132                }
133                if depth > 64 {
134                    return None;
135                }
136            }
137            if depth != 0 {
138                return None;
139            }
140            continue;
141        }
142        if matches!(byte, b'\'' | b'"' | b'`') {
143            index += 1;
144            let mut closed = false;
145            while index < bytes.len() {
146                if bytes[index] == b'\\' {
147                    return None;
148                }
149                let current = bytes[index];
150                index += 1;
151                if current == byte {
152                    if bytes.get(index) == Some(&byte) {
153                        index += 1;
154                        continue;
155                    }
156                    closed = true;
157                    break;
158                }
159            }
160            if !closed {
161                return None;
162            }
163            tokens.push(
164                if byte == b'\'' {
165                    "#value"
166                } else {
167                    "#identifier"
168                }
169                .to_string(),
170            );
171            continue;
172        }
173        if byte.is_ascii_alphabetic() || byte == b'_' {
174            let start = index;
175            index += 1;
176            while index < bytes.len()
177                && (bytes[index].is_ascii_alphanumeric() || bytes[index] == b'_')
178            {
179                index += 1;
180            }
181            tokens.push(sql[start..index].to_ascii_uppercase());
182        } else if byte.is_ascii() {
183            tokens.push(char::from(byte).to_string());
184            index += 1;
185        } else {
186            return None;
187        }
188    }
189    Some(tokens)
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195
196    #[test]
197    fn shared_replica_sql_policy() {
198        let cases: serde_json::Value = serde_json::from_str(include_str!(
199            "../../../testdata/conformance/topology/sql-policy.json"
200        ))
201        .unwrap();
202        for case in cases.as_array().unwrap() {
203            let sql = case["sql"].as_str().unwrap();
204            assert_eq!(
205                sql_route(sql) == SqlRoute::Replica,
206                case["replica_read"].as_bool().unwrap(),
207                "{sql}"
208            );
209        }
210    }
211
212    #[test]
213    fn mutations_locks_and_unknown_functions_stay_on_writer() {
214        for sql in [
215            "UPDATE users SET name = $1",
216            "SELECT * FROM users FOR UPDATE",
217            "SELECT nextval('counter')",
218            "WITH x AS (DELETE FROM users) SELECT * FROM x",
219        ] {
220            assert_eq!(sql_route(sql), SqlRoute::Primary, "{sql}");
221        }
222        assert_eq!(
223            sql_route("/* read */ SELECT COUNT(*) FROM users;"),
224            SqlRoute::Replica
225        );
226        assert_eq!(
227            sql_route("SELECT 'UPDATE; DELETE' AS text"),
228            SqlRoute::Replica
229        );
230        assert_eq!(sql_route("/* transaction */ BEGIN"), SqlRoute::Begin);
231        assert_eq!(sql_route("ROLLBACK TO SAVEPOINT nested"), SqlRoute::Primary);
232        assert_eq!(sql_route("COMMIT"), SqlRoute::End);
233        assert_eq!(sql_route("SELECT 1; BEGIN"), SqlRoute::Session);
234    }
235}
reddb_wire/sql_routing.rs

reddb_wire/
sql_routing.rs
