hawdb_sql_syntax/
lexer.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::{Span, SyntaxError, SyntaxErrorCode, Token, TokenKind};
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub struct LexerConfig {
19    pub max_input_bytes: usize,
20    pub max_tokens: usize,
21    pub max_block_comment_depth: usize,
22}
23
24impl Default for LexerConfig {
25    fn default() -> Self {
26        Self {
27            max_input_bytes: 16 * 1024 * 1024,
28            max_tokens: 1_000_000,
29            max_block_comment_depth: 64,
30        }
31    }
32}
33
34pub fn tokenize(input: &str) -> Result<Vec<Token>, SyntaxError> {
35    tokenize_with_config(input, LexerConfig::default())
36}
37
38/// Reads only the first non-comment token for explicit statement-family dispatch.
39/// The rest of the input is validated by the selected parser. The full-input
40/// byte limit applies to complete tokenization, not to statement dispatch.
41pub fn first_token(input: &str) -> Result<Token, SyntaxError> {
42    Ok(tokenize_inner(input, LexerConfig::default(), true)?[0])
43}
44
45pub fn tokenize_with_config(input: &str, config: LexerConfig) -> Result<Vec<Token>, SyntaxError> {
46    tokenize_inner(input, config, false)
47}
48
49fn tokenize_inner(
50    input: &str,
51    config: LexerConfig,
52    first_only: bool,
53) -> Result<Vec<Token>, SyntaxError> {
54    // Dispatch only reads a prefix. Applying full-input admission here would
55    // also restrict legacy-routed DDL/DML before their parser is selected.
56    if !first_only && input.len() > config.max_input_bytes {
57        return Err(SyntaxError::new(
58            SyntaxErrorCode::InputTooLarge,
59            Span::new(0, input.len()),
60        ));
61    }
62
63    let bytes = input.as_bytes();
64    let mut tokens = Vec::new();
65    let mut position = 0;
66
67    while position < bytes.len() {
68        let byte = bytes[position];
69        if byte.is_ascii_whitespace() {
70            position += 1;
71            continue;
72        }
73
74        if byte == b'-' && bytes.get(position + 1) == Some(&b'-') {
75            position += 2;
76            while position < bytes.len() && !matches!(bytes[position], b'\n' | b'\r') {
77                position += 1;
78            }
79            continue;
80        }
81
82        if byte == b'/' && bytes.get(position + 1) == Some(&b'*') {
83            position = skip_block_comment(input, position, config.max_block_comment_depth)?;
84            continue;
85        }
86
87        let start = position;
88        let kind = match byte {
89            b'u' | b'U'
90                if bytes.get(position + 1) == Some(&b'&')
91                    && bytes.get(position + 2) == Some(&b'\'') =>
92            {
93                position = scan_quoted(
94                    input,
95                    position + 2,
96                    b'\'',
97                    SyntaxErrorCode::UnterminatedString,
98                )?;
99                TokenKind::String
100            }
101            b'e' | b'E' if bytes.get(position + 1) == Some(&b'\'') => {
102                position = scan_escaped_string(input, position)?;
103                TokenKind::String
104            }
105            b'"' => {
106                position = scan_quoted(
107                    input,
108                    position,
109                    b'"',
110                    SyntaxErrorCode::UnterminatedQuotedIdentifier,
111                )?;
112                TokenKind::QuotedIdentifier
113            }
114            b'\'' => {
115                position =
116                    scan_quoted(input, position, b'\'', SyntaxErrorCode::UnterminatedString)?;
117                TokenKind::String
118            }
119            b'$' => {
120                if bytes.get(position + 1).is_some_and(u8::is_ascii_digit) {
121                    let (next, parameter) = scan_parameter(input, position)?;
122                    position = next;
123                    TokenKind::Parameter(parameter)
124                } else if let Some(next) = scan_dollar_quoted_string(input, position)? {
125                    position = next;
126                    TokenKind::String
127                } else {
128                    return Err(SyntaxError::new(
129                        SyntaxErrorCode::InvalidParameter,
130                        Span::new(position, (position + 1).min(input.len())),
131                    )
132                    .expected("a one-based parameter or dollar-quoted string"));
133                }
134            }
135            b'0'..=b'9' => {
136                position = scan_number(input, position);
137                TokenKind::Number
138            }
139            b'(' => single(&mut position, TokenKind::LeftParen),
140            b')' => single(&mut position, TokenKind::RightParen),
141            b'[' => single(&mut position, TokenKind::LeftBracket),
142            b']' => single(&mut position, TokenKind::RightBracket),
143            b'{' => single(&mut position, TokenKind::LeftBrace),
144            b'}' => single(&mut position, TokenKind::RightBrace),
145            b',' => single(&mut position, TokenKind::Comma),
146            b'.' if bytes.get(position + 1).is_some_and(u8::is_ascii_digit) => {
147                position = scan_number(input, position);
148                TokenKind::Number
149            }
150            b'.' => single(&mut position, TokenKind::Dot),
151            b';' => single(&mut position, TokenKind::Semicolon),
152            b'+' => single(&mut position, TokenKind::Plus),
153            b'*' => single(&mut position, TokenKind::Star),
154            b'%' => single(&mut position, TokenKind::Percent),
155            b'&' => single(&mut position, TokenKind::Ampersand),
156            b'^' => single(&mut position, TokenKind::Caret),
157            b'?' => single(&mut position, TokenKind::Question),
158            b':' if bytes.get(position + 1) == Some(&b':') => {
159                position += 2;
160                TokenKind::DoubleColon
161            }
162            b':' => single(&mut position, TokenKind::Colon),
163            b'|' if bytes.get(position + 1) == Some(&b'|') => {
164                position += 2;
165                TokenKind::Concat
166            }
167            b'|' => single(&mut position, TokenKind::Pipe),
168            b'=' => single(&mut position, TokenKind::Equal),
169            b'!' if bytes.get(position + 1) == Some(&b'=') => {
170                position += 2;
171                TokenKind::NotEqual
172            }
173            b'<' if bytes.get(position + 1) == Some(&b'=') => {
174                position += 2;
175                TokenKind::LessOrEqual
176            }
177            b'<' if bytes.get(position + 1) == Some(&b'>') => {
178                position += 2;
179                TokenKind::NotEqual
180            }
181            b'<' if bytes.get(position + 1) == Some(&b'-') => {
182                position += 2;
183                TokenKind::ArrowLeft
184            }
185            b'<' => single(&mut position, TokenKind::Less),
186            b'>' if bytes.get(position + 1) == Some(&b'=') => {
187                position += 2;
188                TokenKind::GreaterOrEqual
189            }
190            b'>' => single(&mut position, TokenKind::Greater),
191            b'-' if bytes.get(position + 1) == Some(&b'>') => {
192                position += 2;
193                TokenKind::ArrowRight
194            }
195            b'-' => single(&mut position, TokenKind::Minus),
196            b'/' => single(&mut position, TokenKind::Slash),
197            _ => {
198                let character = input[position..].chars().next().ok_or_else(|| {
199                    SyntaxError::new(
200                        SyntaxErrorCode::InvalidCharacter,
201                        Span::new(position, position),
202                    )
203                })?;
204                if is_identifier_start(character) {
205                    position = scan_identifier(input, position);
206                    TokenKind::Word
207                } else {
208                    let end = position + character.len_utf8();
209                    return Err(SyntaxError::new(
210                        SyntaxErrorCode::InvalidCharacter,
211                        Span::new(position, end),
212                    )
213                    .found(character.to_string()));
214                }
215            }
216        };
217
218        if kind == TokenKind::Number {
219            for position in start..position {
220                if bytes[position] == b'_'
221                    && (!position
222                        .checked_sub(1)
223                        .is_some_and(|previous| bytes[previous].is_ascii_digit())
224                        || !bytes.get(position + 1).is_some_and(u8::is_ascii_digit))
225                {
226                    return Err(SyntaxError::new(
227                        SyntaxErrorCode::InvalidCharacter,
228                        Span::new(position, position + 1),
229                    )
230                    .expected("a digit on each side of a numeric separator")
231                    .found("_"));
232                }
233            }
234        }
235
236        push_token(
237            &mut tokens,
238            Token::new(kind, Span::new(start, position)),
239            config.max_tokens,
240        )?;
241        if first_only {
242            return Ok(tokens);
243        }
244    }
245
246    tokens.push(Token::new(
247        TokenKind::End,
248        Span::new(input.len(), input.len()),
249    ));
250    Ok(tokens)
251}
252
253fn single(position: &mut usize, kind: TokenKind) -> TokenKind {
254    *position += 1;
255    kind
256}
257
258fn push_token(tokens: &mut Vec<Token>, token: Token, max_tokens: usize) -> Result<(), SyntaxError> {
259    if tokens.len() >= max_tokens {
260        return Err(SyntaxError::new(
261            SyntaxErrorCode::TokenLimitExceeded,
262            token.span,
263        ));
264    }
265    tokens.push(token);
266    Ok(())
267}
268
269fn skip_block_comment(input: &str, start: usize, max_depth: usize) -> Result<usize, SyntaxError> {
270    if max_depth == 0 {
271        return Err(SyntaxError::new(
272            SyntaxErrorCode::CommentNestingLimitExceeded,
273            Span::new(start, (start + 2).min(input.len())),
274        ));
275    }
276    let bytes = input.as_bytes();
277    let mut position = start + 2;
278    let mut depth = 1usize;
279    while position < bytes.len() {
280        match (bytes[position], bytes.get(position + 1).copied()) {
281            (b'/', Some(b'*')) => {
282                depth = depth.checked_add(1).ok_or_else(|| {
283                    SyntaxError::new(
284                        SyntaxErrorCode::CommentNestingLimitExceeded,
285                        Span::new(start, position + 2),
286                    )
287                })?;
288                if depth > max_depth {
289                    return Err(SyntaxError::new(
290                        SyntaxErrorCode::CommentNestingLimitExceeded,
291                        Span::new(start, position + 2),
292                    ));
293                }
294                position += 2;
295            }
296            (b'*', Some(b'/')) => {
297                depth -= 1;
298                position += 2;
299                if depth == 0 {
300                    return Ok(position);
301                }
302            }
303            _ => position += 1,
304        }
305    }
306    Err(SyntaxError::new(
307        SyntaxErrorCode::UnterminatedBlockComment,
308        Span::new(start, input.len()),
309    ))
310}
311
312fn scan_quoted(
313    input: &str,
314    start: usize,
315    quote: u8,
316    error_code: SyntaxErrorCode,
317) -> Result<usize, SyntaxError> {
318    let bytes = input.as_bytes();
319    let mut position = start + 1;
320    while position < bytes.len() {
321        if bytes[position] == quote {
322            if bytes.get(position + 1) == Some(&quote) {
323                position += 2;
324            } else {
325                return Ok(position + 1);
326            }
327        } else {
328            position += 1;
329        }
330    }
331    Err(SyntaxError::new(error_code, Span::new(start, input.len())))
332}
333
334fn scan_escaped_string(input: &str, start: usize) -> Result<usize, SyntaxError> {
335    let bytes = input.as_bytes();
336    let mut position = start + 2;
337    while position < bytes.len() {
338        match bytes[position] {
339            b'\\' => position += 2,
340            b'\'' if bytes.get(position + 1) == Some(&b'\'') => position += 2,
341            b'\'' => return Ok(position + 1),
342            _ => position += 1,
343        }
344    }
345    Err(SyntaxError::new(
346        SyntaxErrorCode::UnterminatedString,
347        Span::new(start, input.len()),
348    ))
349}
350
351fn scan_parameter(input: &str, start: usize) -> Result<(usize, u32), SyntaxError> {
352    let bytes = input.as_bytes();
353    let mut position = start + 1;
354    while bytes.get(position).is_some_and(u8::is_ascii_digit) {
355        position += 1;
356    }
357    let value = input[start + 1..position].parse::<u32>().map_err(|_| {
358        SyntaxError::new(
359            SyntaxErrorCode::InvalidParameter,
360            Span::new(start, position),
361        )
362    })?;
363    if value == 0 {
364        return Err(SyntaxError::new(
365            SyntaxErrorCode::InvalidParameter,
366            Span::new(start, position),
367        )
368        .expected("a one-based PostgreSQL parameter"));
369    }
370    Ok((position, value))
371}
372
373fn scan_dollar_quoted_string(input: &str, start: usize) -> Result<Option<usize>, SyntaxError> {
374    let bytes = input.as_bytes();
375    let mut delimiter_end = start + 1;
376    while let Some(byte) = bytes.get(delimiter_end) {
377        if *byte == b'$' {
378            let delimiter = &input[start..=delimiter_end];
379            let content_start = delimiter_end + 1;
380            return input[content_start..]
381                .find(delimiter)
382                .map(|offset| Some(content_start + offset + delimiter.len()))
383                .ok_or_else(|| {
384                    SyntaxError::new(
385                        SyntaxErrorCode::UnterminatedDollarQuotedString,
386                        Span::new(start, input.len()),
387                    )
388                });
389        }
390        if !byte.is_ascii_alphanumeric() && *byte != b'_' {
391            return Ok(None);
392        }
393        delimiter_end += 1;
394    }
395    Ok(None)
396}
397
398fn scan_number(input: &str, start: usize) -> usize {
399    let bytes = input.as_bytes();
400    let mut position = start;
401    while bytes
402        .get(position)
403        .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
404    {
405        position += 1;
406    }
407    if bytes.get(position) == Some(&b'.') && bytes.get(position + 1) != Some(&b'.') {
408        position += 1;
409        while bytes
410            .get(position)
411            .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
412        {
413            position += 1;
414        }
415    }
416    if matches!(bytes.get(position), Some(b'e' | b'E')) {
417        let exponent = position;
418        position += 1;
419        if matches!(bytes.get(position), Some(b'+' | b'-')) {
420            position += 1;
421        }
422        let digit_start = position;
423        while bytes
424            .get(position)
425            .is_some_and(|byte| byte.is_ascii_digit() || *byte == b'_')
426        {
427            position += 1;
428        }
429        if digit_start == position {
430            return exponent;
431        }
432    }
433    position
434}
435
436fn scan_identifier(input: &str, start: usize) -> usize {
437    let mut end = start;
438    for (offset, character) in input[start..].char_indices() {
439        if offset == 0 {
440            if !is_identifier_start(character) {
441                break;
442            }
443        } else if !is_identifier_continue(character) {
444            break;
445        }
446        end = start + offset + character.len_utf8();
447    }
448    end
449}
450
451fn is_identifier_start(character: char) -> bool {
452    character == '_' || character.is_alphabetic()
453}
454
455fn is_identifier_continue(character: char) -> bool {
456    is_identifier_start(character) || character.is_numeric() || character == '$'
457}
458
459#[cfg(test)]
460mod tests {
461    use super::*;
462
463    fn kinds(input: &str) -> Vec<TokenKind> {
464        tokenize(input)
465            .expect("valid SQL")
466            .into_iter()
467            .map(|token| token.kind)
468            .collect()
469    }
470
471    #[test]
472    fn tokenizes_postgres_sql_pgq_surface() {
473        assert_eq!(
474            kinds(
475                "SELECT p.name FROM GRAPH_TABLE (g MATCH (p IS person)-[e IS knows]->(q) \
476                 COLUMNS (p.name AS name, $1)) AS gt;",
477            ),
478            vec![
479                TokenKind::Word,
480                TokenKind::Word,
481                TokenKind::Dot,
482                TokenKind::Word,
483                TokenKind::Word,
484                TokenKind::Word,
485                TokenKind::LeftParen,
486                TokenKind::Word,
487                TokenKind::Word,
488                TokenKind::LeftParen,
489                TokenKind::Word,
490                TokenKind::Word,
491                TokenKind::Word,
492                TokenKind::RightParen,
493                TokenKind::Minus,
494                TokenKind::LeftBracket,
495                TokenKind::Word,
496                TokenKind::Word,
497                TokenKind::Word,
498                TokenKind::RightBracket,
499                TokenKind::ArrowRight,
500                TokenKind::LeftParen,
501                TokenKind::Word,
502                TokenKind::RightParen,
503                TokenKind::Word,
504                TokenKind::LeftParen,
505                TokenKind::Word,
506                TokenKind::Dot,
507                TokenKind::Word,
508                TokenKind::Word,
509                TokenKind::Word,
510                TokenKind::Comma,
511                TokenKind::Parameter(1),
512                TokenKind::RightParen,
513                TokenKind::RightParen,
514                TokenKind::Word,
515                TokenKind::Word,
516                TokenKind::Semicolon,
517                TokenKind::End,
518            ]
519        );
520    }
521
522    #[test]
523    fn skips_nested_comments_and_preserves_byte_spans() {
524        let input = "/* outer /* nested */ done */ SELECT \"na\"\"me\"";
525        let tokens = tokenize(input).expect("valid nested comment");
526        assert_eq!(tokens[0].text(input), "SELECT");
527        assert_eq!(tokens[1].kind, TokenKind::QuotedIdentifier);
528        assert_eq!(tokens[1].text(input), "\"na\"\"me\"");
529    }
530
531    #[test]
532    fn recognizes_dollar_quoted_strings_and_unicode_identifiers() {
533        let input = "SELECT 数据, $tag$MATCH (n)$tag$";
534        let tokens = tokenize(input).expect("valid PostgreSQL tokens");
535        assert_eq!(tokens[1].text(input), "数据");
536        assert_eq!(tokens[3].kind, TokenKind::String);
537    }
538
539    #[test]
540    fn rejects_zero_parameter_and_unterminated_input() {
541        let error = tokenize("SELECT $0").expect_err("zero parameter must fail");
542        assert_eq!(error.code, SyntaxErrorCode::InvalidParameter);
543
544        let error = tokenize("SELECT /*").expect_err("comment must terminate");
545        assert_eq!(error.code, SyntaxErrorCode::UnterminatedBlockComment);
546    }
547
548    #[test]
549    fn enforces_token_and_comment_depth_limits() {
550        let error = tokenize_with_config(
551            "SELECT one",
552            LexerConfig {
553                max_tokens: 1,
554                ..LexerConfig::default()
555            },
556        )
557        .expect_err("token budget must fail");
558        assert_eq!(error.code, SyntaxErrorCode::TokenLimitExceeded);
559
560        let error = tokenize_with_config(
561            "/* /* nested */ */",
562            LexerConfig {
563                max_block_comment_depth: 1,
564                ..LexerConfig::default()
565            },
566        )
567        .expect_err("comment depth must fail");
568        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
569
570        let error = tokenize_with_config(
571            "/**/",
572            LexerConfig {
573                max_block_comment_depth: 0,
574                ..LexerConfig::default()
575            },
576        )
577        .expect_err("zero comment budget must fail");
578        assert_eq!(error.code, SyntaxErrorCode::CommentNestingLimitExceeded);
579    }
580}
hawdb_sql_syntax/lexer.rs

hawdb_sql_syntax/
lexer.rs
