Skip to main content

sofie_core/
sql_lexer.rs

1//! Postgres's lexical rules where they hide text: `--` and nested `/* */`
2//! comments, `''` and `E'\''` escapes, `"quoted"` names and `$tag$` bodies.
3
4#[derive(Clone, Debug, PartialEq, Eq)]
5pub enum Token {
6    Word(String),
7    /// A `"quoted"` identifier, never a keyword: `ADD COLUMN "check"` adds a column.
8    Quoted(String),
9    Text(String),
10    Symbol(char),
11}
12
13impl Token {
14    #[must_use]
15    pub fn keyword(&self) -> &str {
16        match self {
17            Self::Word(word) => word,
18            _ => "",
19        }
20    }
21}
22
23#[derive(Debug, Default)]
24pub struct Lexed {
25    pub tokens: Vec<Token>,
26    pub lines: Vec<usize>,
27    pub comments: Vec<(usize, String)>,
28}
29
30/// A `$tag$` body is lexed on its own: its words stay visible, its `;` ends no statement.
31#[must_use]
32pub fn lex(sql: &str) -> Lexed {
33    let characters: Vec<char> = sql.chars().collect();
34    let mut lexed = Lexed::default();
35    let mut line = 1;
36    let mut index = 0;
37    while index < characters.len() {
38        let rest = &characters[index..];
39        let (token, length) = match rest[0] {
40            '-' if rest.get(1) == Some(&'-') => {
41                let length = rest
42                    .iter()
43                    .position(|next| *next == '\n')
44                    .unwrap_or(rest.len());
45                let comment: String = rest[2..length].iter().collect();
46                lexed.comments.push((line, comment.trim().to_owned()));
47                (None, length)
48            }
49            '/' if rest.get(1) == Some(&'*') => (None, block_comment_length(rest)),
50            '$' if index == 0 || !is_word_character(characters[index - 1]) => {
51                match dollar_tag_length(rest) {
52                    Some(tag) => {
53                        let (body, length) = dollar_quoted(rest, tag);
54                        lexed.append_body(lex(&body), line - 1);
55                        (None, length)
56                    }
57                    None => (Some(Token::Symbol('$')), 1),
58                }
59            }
60            '$' => (Some(Token::Symbol('$')), 1),
61            quote @ ('\'' | '"') => {
62                let escapes = quote == '\'' && is_escape_prefix(&lexed, &characters, index);
63                if escapes {
64                    lexed.tokens.pop();
65                    lexed.lines.pop();
66                }
67                let (content, length) = quoted(rest, escapes);
68                let token = if quote == '"' {
69                    Token::Quoted(content)
70                } else {
71                    Token::Text(content)
72                };
73                (Some(token), length)
74            }
75            first if is_word_character(first) => {
76                let length = rest
77                    .iter()
78                    .position(|next| !is_word_character(*next) && *next != '$')
79                    .unwrap_or(rest.len());
80                let word = rest[..length].iter().flat_map(|next| next.to_lowercase());
81                (Some(Token::Word(word.collect())), length)
82            }
83            whitespace if whitespace.is_whitespace() => (None, 1),
84            symbol => (Some(Token::Symbol(symbol)), 1),
85        };
86        if let Some(token) = token {
87            lexed.tokens.push(token);
88            lexed.lines.push(line);
89        }
90        line += rest[..length].iter().filter(|next| **next == '\n').count();
91        index += length;
92    }
93    lexed
94}
95
96impl Lexed {
97    fn append_body(&mut self, body: Lexed, lines_before: usize) {
98        for (token, line) in body.tokens.into_iter().zip(body.lines) {
99            if token != Token::Symbol(';') {
100                self.tokens.push(token);
101                self.lines.push(line + lines_before);
102            }
103        }
104        let comments = body.comments.into_iter();
105        self.comments
106            .extend(comments.map(|(line, comment)| (line + lines_before, comment)));
107    }
108}
109
110/// A `$` also continues a word, as `t$x$` names one table; it never starts one.
111fn is_word_character(character: char) -> bool {
112    character.is_alphanumeric() || character == '_'
113}
114
115fn block_comment_length(rest: &[char]) -> usize {
116    let mut depth = 0;
117    let mut index = 0;
118    while index + 1 < rest.len() {
119        match (rest[index], rest[index + 1]) {
120            ('/', '*') => depth += 1,
121            ('*', '/') => depth -= 1,
122            _ => {
123                index += 1;
124                continue;
125            }
126        }
127        index += 2;
128        if depth == 0 {
129            return index;
130        }
131    }
132    rest.len()
133}
134
135/// `$$` or `$tag$`; a `$1` parameter is no tag.
136fn dollar_tag_length(rest: &[char]) -> Option<usize> {
137    let close = rest[1..].iter().position(|next| *next == '$')? + 1;
138    let tag = &rest[1..close];
139    let valid = tag.first().is_none_or(|first| !first.is_ascii_digit())
140        && tag
141            .iter()
142            .all(|next| next.is_alphanumeric() || *next == '_');
143    valid.then_some(close + 1)
144}
145
146fn dollar_quoted(rest: &[char], tag: usize) -> (String, usize) {
147    let delimiter = &rest[..tag];
148    let body = &rest[tag..];
149    match body.windows(tag).position(|window| window == delimiter) {
150        Some(end) => (body[..end].iter().collect(), tag + end + tag),
151        None => (body.iter().collect(), rest.len()),
152    }
153}
154
155fn is_escape_prefix(lexed: &Lexed, characters: &[char], index: usize) -> bool {
156    index > 0
157        && matches!(characters[index - 1], 'e' | 'E')
158        && matches!(lexed.tokens.last(), Some(Token::Word(word)) if word == "e")
159}
160
161fn quoted(rest: &[char], escapes: bool) -> (String, usize) {
162    let quote = rest[0];
163    let mut content = String::new();
164    let mut index = 1;
165    while let Some(&next) = rest.get(index) {
166        match next {
167            '\\' if escapes => {
168                content.extend(rest.get(index + 1));
169                index += 2;
170            }
171            _ if next == quote && rest.get(index + 1) == Some(&quote) => {
172                content.push(quote);
173                index += 2;
174            }
175            _ if next == quote => return (content, index + 1),
176            _ => {
177                content.push(next);
178                index += 1;
179            }
180        }
181    }
182    (content, rest.len())
183}
184
185#[cfg(test)]
186mod tests {
187    use super::{Token, lex};
188
189    fn words(sql: &str) -> Vec<String> {
190        lex(sql)
191            .tokens
192            .iter()
193            .filter_map(|token| match token {
194                Token::Word(word) => Some(word.clone()),
195                _ => None,
196            })
197            .collect()
198    }
199
200    #[test]
201    fn block_comments_dollar_bodies_and_escapes_hide_no_statement() {
202        let lexed = lex(
203            "/* don't run /* nested */ at peak; */\nCREATE INDEX a ON t (x);\n\
204             COMMENT ON TABLE t IS $c$ don't; $c$;\nSELECT E'it\\'s', 'it''s', $1;\n\
205             CREATE INDEX b ON t (y);",
206        );
207
208        let texts: Vec<&str> = lexed
209            .tokens
210            .iter()
211            .filter_map(|token| match token {
212                Token::Text(text) => Some(text.as_str()),
213                _ => None,
214            })
215            .collect();
216        assert_eq!(texts, ["t; ", "it's", "it's"]);
217        let create = |name: &str| {
218            let word = Token::Word(name.to_owned());
219            let at = lexed.tokens.iter().position(|token| *token == word);
220            lexed.lines[at.unwrap()]
221        };
222        assert_eq!((create("a"), create("b")), (2, 5));
223        let statements = lexed.tokens.split(|token| *token == Token::Symbol(';'));
224        assert_eq!(statements.count(), 5);
225        assert_eq!(
226            words("/* ' */ CREATE INDEX a ON t (x);"),
227            ["create", "index", "a", "on", "t", "x"]
228        );
229    }
230
231    fn texts(sql: &str) -> Vec<String> {
232        lex(sql)
233            .tokens
234            .iter()
235            .filter_map(|token| match token {
236                Token::Text(text) => Some(text.clone()),
237                _ => None,
238            })
239            .collect()
240    }
241
242    #[test]
243    fn a_dollar_inside_a_word_opens_no_body() {
244        assert_eq!(
245            words("CREATE TABLE t$x$ (a int); CREATE INDEX b ON t (y);"),
246            [
247                "create", "table", "t$x$", "a", "int", "create", "index", "b", "on", "t", "y"
248            ]
249        );
250        assert_eq!(words("SELECT 'a'$q$b$q$"), ["select", "b"]);
251    }
252
253    #[test]
254    fn an_unterminated_quote_or_comment_ends_with_the_text() {
255        assert_eq!(texts("SELECT 'abc"), ["abc"]);
256        assert_eq!(
257            words("SELECT 1 /* never closed; CREATE INDEX"),
258            ["select", "1"]
259        );
260        assert_eq!(
261            words("SELECT $a$ never closed; CREATE INDEX"),
262            ["select", "never", "closed", "create", "index"]
263        );
264    }
265
266    #[test]
267    fn unicode_and_backslash_escaped_strings_end_where_postgres_ends_them() {
268        let unicode = lex("SELECT U&'d\\0061t\\'; CREATE INDEX b ON t (y)");
269        let backslash = lex("SELECT E'\\\\'; CREATE INDEX b ON t (y)");
270
271        for lexed in [&unicode, &backslash] {
272            let statements = lexed.tokens.split(|token| *token == Token::Symbol(';'));
273            assert_eq!(statements.count(), 2, "{lexed:?}");
274        }
275        assert!(
276            unicode
277                .tokens
278                .contains(&Token::Text("d\\0061t\\".to_owned()))
279        );
280        assert!(backslash.tokens.contains(&Token::Text("\\".to_owned())));
281    }
282
283    #[test]
284    fn crlf_lines_are_counted_once() {
285        let lexed = lex("SELECT 1;\r\n-- note\r\nCREATE INDEX b ON t (y);\r\n");
286
287        let create = lexed
288            .tokens
289            .iter()
290            .position(|token| *token == Token::Word("create".to_owned()))
291            .unwrap();
292        assert_eq!(lexed.lines[create], 3);
293        assert_eq!(lexed.comments, [(2, "note".to_owned())]);
294    }
295
296    #[test]
297    fn quoted_names_and_line_comments_hold_no_word() {
298        let lexed = lex("ALTER TABLE t ADD \"Check\" text; -- Not now; later\nSELECT 1");
299
300        assert!(lexed.tokens.contains(&Token::Quoted("Check".to_owned())));
301        assert_eq!(lexed.comments, [(1, "Not now; later".to_owned())]);
302        assert_eq!(
303            words("ALTER TABLE t ADD \"Check\" text; -- Not now\nSELECT 1"),
304            ["alter", "table", "t", "add", "text", "select", "1"]
305        );
306    }
307}