1#[derive(Clone, Debug, PartialEq, Eq)]
5pub enum Token {
6 Word(String),
7 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#[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
110fn 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
135fn 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("e) => {
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}