sql_cli/sql/parser/
lexer.rs

1//! SQL Lexer - Tokenization of SQL queries
2//!
3//! This module handles the conversion of raw SQL text into tokens
4//! that can be consumed by the parser.
5
6/// Lexer mode - controls whether comments are preserved or skipped
7#[derive(Debug, Clone, Copy, PartialEq)]
8pub enum LexerMode {
9    /// Standard mode - skip comments (current default behavior)
10    SkipComments,
11    /// Preserve mode - tokenize comments as tokens
12    PreserveComments,
13}
14
15impl Default for LexerMode {
16    fn default() -> Self {
17        LexerMode::SkipComments
18    }
19}
20
21#[derive(Debug, Clone, PartialEq)]
22pub enum Token {
23    // Keywords
24    Select,
25    From,
26    Where,
27    With, // WITH clause for CTEs
28    And,
29    Or,
30    In,
31    Not,
32    Between,
33    Like,
34    ILike, // Case-insensitive LIKE (PostgreSQL)
35    Is,
36    Null,
37    OrderBy,
38    GroupBy,
39    Having,
40    Qualify,
41    As,
42    Asc,
43    Desc,
44    Limit,
45    Offset,
46    Into,      // INTO keyword for temporary tables
47    DateTime,  // DateTime constructor
48    Case,      // CASE expression
49    When,      // WHEN clause
50    Then,      // THEN clause
51    Else,      // ELSE clause
52    End,       // END keyword
53    Distinct,  // DISTINCT keyword for aggregate functions
54    Over,      // OVER keyword for window functions
55    Partition, // PARTITION keyword for window functions
56    By,        // BY keyword (used with PARTITION BY, ORDER BY)
57    Exclude,   // EXCLUDE keyword (for SELECT * EXCLUDE)
58    // Note: REPLACE is NOT a keyword - it's handled as a function name
59    // to avoid conflicting with the REPLACE() string function
60
61    // PIVOT/UNPIVOT keywords
62    Pivot,   // PIVOT keyword for row-to-column transformation
63    Unpivot, // UNPIVOT keyword for column-to-row transformation
64    For,     // FOR keyword (used in PIVOT: FOR column IN (...))
65
66    // Window frame keywords
67    Rows,      // ROWS frame type
68    Range,     // RANGE frame type
69    Unbounded, // UNBOUNDED for frame bounds
70    Preceding, // PRECEDING for frame bounds
71    Following, // FOLLOWING for frame bounds
72    Current,   // CURRENT for CURRENT ROW
73    Row,       // ROW for CURRENT ROW
74
75    // Set operation keywords
76    Union,     // UNION
77    Intersect, // INTERSECT
78    Except,    // EXCEPT
79
80    // Special CTE keywords
81    Web,  // WEB (for WEB CTEs)
82    File, // FILE (for FILE CTEs — filesystem metadata)
83
84    // Row expansion functions
85    Unnest, // UNNEST (for expanding delimited strings into rows)
86
87    // JOIN keywords
88    Join,  // JOIN keyword
89    Inner, // INNER JOIN
90    Left,  // LEFT JOIN
91    Right, // RIGHT JOIN
92    Full,  // FULL JOIN
93    Outer, // OUTER keyword (LEFT OUTER, RIGHT OUTER, FULL OUTER)
94    On,    // ON keyword for join conditions
95    Cross, // CROSS JOIN
96
97    // Literals
98    Identifier(String),
99    QuotedIdentifier(String), // For "Customer Id" style identifiers
100    StringLiteral(String),
101    JsonBlock(String), // For $JSON$...$ JSON$ delimited blocks
102    NumberLiteral(String),
103    Star,
104
105    // Operators
106    Dot,
107    Comma,
108    /// Statement terminator. Previously fell through the catch-all and lexed as
109    /// `Identifier(";")`, which the parser then silently ignored along with
110    /// everything after it (P13).
111    Semicolon,
112    Colon,
113    LeftParen,
114    RightParen,
115    Equal,
116    NotEqual,
117    LessThan,
118    GreaterThan,
119    LessThanOrEqual,
120    GreaterThanOrEqual,
121
122    // Arithmetic operators
123    Plus,
124    Minus,
125    Divide,
126    Modulo,
127
128    // String operators
129    Concat, // || for string concatenation
130
131    // Comments (preserved for formatting)
132    LineComment(String),  // -- comment text (without the -- prefix)
133    BlockComment(String), // /* comment text */ (without delimiters)
134
135    // Special
136    Eof,
137}
138
139impl Token {
140    /// Check if a string is a SQL keyword and return corresponding token
141    pub fn from_keyword(s: &str) -> Option<Token> {
142        match s.to_uppercase().as_str() {
143            "SELECT" => Some(Token::Select),
144            "FROM" => Some(Token::From),
145            "WHERE" => Some(Token::Where),
146            "WITH" => Some(Token::With),
147            "AND" => Some(Token::And),
148            "OR" => Some(Token::Or),
149            "IN" => Some(Token::In),
150            "NOT" => Some(Token::Not),
151            "BETWEEN" => Some(Token::Between),
152            "LIKE" => Some(Token::Like),
153            "ILIKE" => Some(Token::ILike),
154            "IS" => Some(Token::Is),
155            "NULL" => Some(Token::Null),
156            "ORDER" => Some(Token::OrderBy),
157            "GROUP" => Some(Token::GroupBy),
158            "HAVING" => Some(Token::Having),
159            "QUALIFY" => Some(Token::Qualify),
160            "AS" => Some(Token::As),
161            "ASC" => Some(Token::Asc),
162            "DESC" => Some(Token::Desc),
163            "LIMIT" => Some(Token::Limit),
164            "OFFSET" => Some(Token::Offset),
165            "INTO" => Some(Token::Into),
166            "DISTINCT" => Some(Token::Distinct),
167            "EXCLUDE" => Some(Token::Exclude),
168            "PIVOT" => Some(Token::Pivot),
169            "UNPIVOT" => Some(Token::Unpivot),
170            "FOR" => Some(Token::For),
171            "CASE" => Some(Token::Case),
172            "WHEN" => Some(Token::When),
173            "THEN" => Some(Token::Then),
174            "ELSE" => Some(Token::Else),
175            "END" => Some(Token::End),
176            "OVER" => Some(Token::Over),
177            "PARTITION" => Some(Token::Partition),
178            "BY" => Some(Token::By),
179            "ROWS" => Some(Token::Rows),
180            "RANGE" => Some(Token::Range),
181            "UNBOUNDED" => Some(Token::Unbounded),
182            "PRECEDING" => Some(Token::Preceding),
183            "FOLLOWING" => Some(Token::Following),
184            "CURRENT" => Some(Token::Current),
185            "ROW" => Some(Token::Row),
186            "UNION" => Some(Token::Union),
187            "INTERSECT" => Some(Token::Intersect),
188            "EXCEPT" => Some(Token::Except),
189            "WEB" => Some(Token::Web),
190            "FILE" => Some(Token::File),
191            "UNNEST" => Some(Token::Unnest),
192            "JOIN" => Some(Token::Join),
193            "INNER" => Some(Token::Inner),
194            "LEFT" => Some(Token::Left),
195            "RIGHT" => Some(Token::Right),
196            "FULL" => Some(Token::Full),
197            "OUTER" => Some(Token::Outer),
198            "ON" => Some(Token::On),
199            "CROSS" => Some(Token::Cross),
200            _ => None,
201        }
202    }
203
204    /// Check if token is a logical operator
205    pub fn is_logical_operator(&self) -> bool {
206        matches!(self, Token::And | Token::Or)
207    }
208
209    /// Check if token is a join type
210    pub fn is_join_type(&self) -> bool {
211        matches!(
212            self,
213            Token::Inner | Token::Left | Token::Right | Token::Full | Token::Cross
214        )
215    }
216
217    /// Check if token ends a clause
218    pub fn is_clause_terminator(&self) -> bool {
219        matches!(
220            self,
221            Token::OrderBy
222                | Token::GroupBy
223                | Token::Having
224                | Token::Limit
225                | Token::Offset
226                | Token::Union
227                | Token::Intersect
228                | Token::Except
229        )
230    }
231
232    /// Get the string representation of a keyword token
233    /// Returns the keyword as it would appear in SQL (uppercase)
234    pub fn as_keyword_str(&self) -> Option<&'static str> {
235        match self {
236            Token::Select => Some("SELECT"),
237            Token::From => Some("FROM"),
238            Token::Where => Some("WHERE"),
239            Token::With => Some("WITH"),
240            Token::And => Some("AND"),
241            Token::Or => Some("OR"),
242            Token::In => Some("IN"),
243            Token::Not => Some("NOT"),
244            Token::Between => Some("BETWEEN"),
245            Token::Like => Some("LIKE"),
246            Token::ILike => Some("ILIKE"),
247            Token::Is => Some("IS"),
248            Token::Null => Some("NULL"),
249            Token::OrderBy => Some("ORDER BY"),
250            Token::GroupBy => Some("GROUP BY"),
251            Token::Having => Some("HAVING"),
252            Token::Qualify => Some("QUALIFY"),
253            Token::As => Some("AS"),
254            Token::Asc => Some("ASC"),
255            Token::Desc => Some("DESC"),
256            Token::Limit => Some("LIMIT"),
257            Token::Offset => Some("OFFSET"),
258            Token::Into => Some("INTO"),
259            Token::Distinct => Some("DISTINCT"),
260            Token::Exclude => Some("EXCLUDE"),
261            Token::Pivot => Some("PIVOT"),
262            Token::Unpivot => Some("UNPIVOT"),
263            Token::For => Some("FOR"),
264            Token::Case => Some("CASE"),
265            Token::When => Some("WHEN"),
266            Token::Then => Some("THEN"),
267            Token::Else => Some("ELSE"),
268            Token::End => Some("END"),
269            Token::Join => Some("JOIN"),
270            Token::Inner => Some("INNER"),
271            Token::Left => Some("LEFT"),
272            Token::Right => Some("RIGHT"),
273            Token::Full => Some("FULL"),
274            Token::Cross => Some("CROSS"),
275            Token::On => Some("ON"),
276            Token::Union => Some("UNION"),
277            Token::Intersect => Some("INTERSECT"),
278            Token::Except => Some("EXCEPT"),
279            Token::Over => Some("OVER"),
280            Token::Partition => Some("PARTITION"),
281            Token::By => Some("BY"),
282            Token::Rows => Some("ROWS"),
283            Token::Range => Some("RANGE"),
284            Token::Preceding => Some("PRECEDING"),
285            Token::Following => Some("FOLLOWING"),
286            Token::Current => Some("CURRENT"),
287            Token::Row => Some("ROW"),
288            Token::Unbounded => Some("UNBOUNDED"),
289            Token::DateTime => Some("DATETIME"),
290            _ => None,
291        }
292    }
293}
294
295#[derive(Debug, Clone)]
296pub struct Lexer {
297    input: Vec<char>,
298    position: usize,
299    current_char: Option<char>,
300    mode: LexerMode,
301}
302
303impl Lexer {
304    #[must_use]
305    pub fn new(input: &str) -> Self {
306        Self::with_mode(input, LexerMode::default())
307    }
308
309    /// Create a new lexer with specified mode
310    #[must_use]
311    pub fn with_mode(input: &str, mode: LexerMode) -> Self {
312        let chars: Vec<char> = input.chars().collect();
313        let current = chars.first().copied();
314        Self {
315            input: chars,
316            position: 0,
317            current_char: current,
318            mode,
319        }
320    }
321
322    fn advance(&mut self) {
323        self.position += 1;
324        self.current_char = self.input.get(self.position).copied();
325    }
326
327    fn peek(&self, offset: usize) -> Option<char> {
328        self.input.get(self.position + offset).copied()
329    }
330
331    /// Peek ahead n characters and return as a string
332    fn peek_string(&self, n: usize) -> String {
333        let mut result = String::new();
334        for i in 0..n {
335            if let Some(ch) = self.input.get(self.position + i) {
336                result.push(*ch);
337            } else {
338                break;
339            }
340        }
341        result
342    }
343
344    /// Read a JSON block delimited by $JSON$...$JSON$
345    /// Consumes the opening delimiter and reads until closing $JSON$
346    fn read_json_block(&mut self) -> String {
347        let mut result = String::new();
348
349        // Skip opening $JSON$
350        for _ in 0..6 {
351            self.advance();
352        }
353
354        // Read until we find closing $JSON$
355        while let Some(ch) = self.current_char {
356            // Check if we're at the closing delimiter
357            if ch == '$' && self.peek_string(6) == "$JSON$" {
358                // Skip closing $JSON$
359                for _ in 0..6 {
360                    self.advance();
361                }
362                break;
363            }
364            result.push(ch);
365            self.advance();
366        }
367
368        result
369    }
370
371    fn skip_whitespace(&mut self) {
372        while let Some(ch) = self.current_char {
373            if ch.is_whitespace() {
374                self.advance();
375            } else {
376                break;
377            }
378        }
379    }
380
381    /// Read a line comment and return its content (without the -- prefix)
382    fn read_line_comment(&mut self) -> String {
383        let mut result = String::new();
384
385        // Skip '--'
386        self.advance();
387        self.advance();
388
389        // Read until end of line or EOF
390        while let Some(ch) = self.current_char {
391            if ch == '\n' {
392                self.advance(); // consume the newline
393                break;
394            }
395            result.push(ch);
396            self.advance();
397        }
398
399        result
400    }
401
402    /// Read a block comment and return its content (without /* */ delimiters)
403    fn read_block_comment(&mut self) -> String {
404        let mut result = String::new();
405
406        // Skip '/*'
407        self.advance();
408        self.advance();
409
410        // Read until we find '*/'
411        while let Some(ch) = self.current_char {
412            if ch == '*' && self.peek(1) == Some('/') {
413                self.advance(); // skip '*'
414                self.advance(); // skip '/'
415                break;
416            }
417            result.push(ch);
418            self.advance();
419        }
420
421        result
422    }
423
424    /// Skip whitespace and comments (for backwards compatibility with parser)
425    /// This is the old behavior that discards comments
426    fn skip_whitespace_and_comments(&mut self) {
427        loop {
428            // Skip whitespace
429            while let Some(ch) = self.current_char {
430                if ch.is_whitespace() {
431                    self.advance();
432                } else {
433                    break;
434                }
435            }
436
437            // Check for comments
438            match self.current_char {
439                Some('-') if self.peek(1) == Some('-') => {
440                    // Single-line comment: skip until end of line
441                    self.advance(); // skip first '-'
442                    self.advance(); // skip second '-'
443                    while let Some(ch) = self.current_char {
444                        self.advance();
445                        if ch == '\n' {
446                            break;
447                        }
448                    }
449                }
450                Some('/') if self.peek(1) == Some('*') => {
451                    // Multi-line comment: skip until */
452                    self.advance(); // skip '/'
453                    self.advance(); // skip '*'
454                    while let Some(ch) = self.current_char {
455                        if ch == '*' && self.peek(1) == Some('/') {
456                            self.advance(); // skip '*'
457                            self.advance(); // skip '/'
458                            break;
459                        }
460                        self.advance();
461                    }
462                }
463                _ => {
464                    // No more comments or whitespace
465                    break;
466                }
467            }
468        }
469    }
470
471    fn read_identifier(&mut self) -> String {
472        let mut result = String::new();
473        while let Some(ch) = self.current_char {
474            if ch.is_alphanumeric() || ch == '_' {
475                result.push(ch);
476                self.advance();
477            } else {
478                break;
479            }
480        }
481        result
482    }
483
484    fn read_string(&mut self) -> String {
485        let mut result = String::new();
486        let quote_char = self.current_char.unwrap(); // ' or "
487        self.advance(); // skip opening quote
488
489        while let Some(ch) = self.current_char {
490            if ch == quote_char {
491                // SQL escapes a quote by doubling it: 'O''Brien' is one literal
492                // meaning O'Brien, and "a""b" likewise for quoted identifiers.
493                // Without this the literal ended at the first inner quote and the
494                // remainder became a SECOND literal — which the parser then
495                // silently discarded along with the rest of the statement (P13),
496                // turning `WHERE name = 'O''Brien'` into `WHERE name = 'O'`.
497                if self.peek(1) == Some(quote_char) {
498                    result.push(quote_char);
499                    self.advance(); // consume the first quote
500                    self.advance(); // consume the second
501                    continue;
502                }
503                self.advance(); // skip closing quote
504                break;
505            }
506            result.push(ch);
507            self.advance();
508        }
509        result
510    }
511
512    fn read_number(&mut self) -> String {
513        let mut result = String::new();
514        let has_e = false;
515
516        // Read the main number part (including decimal point)
517        while let Some(ch) = self.current_char {
518            if !has_e && (ch.is_numeric() || ch == '.') {
519                result.push(ch);
520                self.advance();
521            } else if (ch == 'e' || ch == 'E') && !has_e && !result.is_empty() {
522                // Handle scientific notation
523                result.push(ch);
524                self.advance();
525                let _ = has_e; // We don't allow multiple 'e' characters, so break after this
526
527                // Check for optional sign after 'e'
528                if let Some(sign) = self.current_char {
529                    if sign == '+' || sign == '-' {
530                        result.push(sign);
531                        self.advance();
532                    }
533                }
534
535                // Read exponent digits
536                while let Some(digit) = self.current_char {
537                    if digit.is_numeric() {
538                        result.push(digit);
539                        self.advance();
540                    } else {
541                        break;
542                    }
543                }
544                break; // Done reading the number
545            } else {
546                break;
547            }
548        }
549        result
550    }
551
552    /// Get next token while preserving comments as tokens
553    /// This is the new behavior for comment-aware formatting
554    pub fn next_token_with_comments(&mut self) -> Token {
555        // Only skip whitespace, NOT comments
556        self.skip_whitespace();
557
558        match self.current_char {
559            None => Token::Eof,
560            // Handle comments as tokens
561            Some('-') if self.peek(1) == Some('-') => {
562                let comment_text = self.read_line_comment();
563                Token::LineComment(comment_text)
564            }
565            Some('/') if self.peek(1) == Some('*') => {
566                let comment_text = self.read_block_comment();
567                Token::BlockComment(comment_text)
568            }
569            Some('*') => {
570                self.advance();
571                Token::Star
572            }
573            Some('+') => {
574                self.advance();
575                Token::Plus
576            }
577            Some('/') => {
578                // Regular division (comment case handled above)
579                self.advance();
580                Token::Divide
581            }
582            Some('%') => {
583                self.advance();
584                Token::Modulo
585            }
586            Some('.') => {
587                self.advance();
588                Token::Dot
589            }
590            Some(',') => {
591                self.advance();
592                Token::Comma
593            }
594            Some(';') => {
595                self.advance();
596                Token::Semicolon
597            }
598            Some(':') => {
599                self.advance();
600                Token::Colon
601            }
602            Some('(') => {
603                self.advance();
604                Token::LeftParen
605            }
606            Some(')') => {
607                self.advance();
608                Token::RightParen
609            }
610            Some('=') => {
611                self.advance();
612                Token::Equal
613            }
614            Some('<') => {
615                self.advance();
616                if self.current_char == Some('=') {
617                    self.advance();
618                    Token::LessThanOrEqual
619                } else if self.current_char == Some('>') {
620                    self.advance();
621                    Token::NotEqual
622                } else {
623                    Token::LessThan
624                }
625            }
626            Some('>') => {
627                self.advance();
628                if self.current_char == Some('=') {
629                    self.advance();
630                    Token::GreaterThanOrEqual
631                } else {
632                    Token::GreaterThan
633                }
634            }
635            Some('!') if self.peek(1) == Some('=') => {
636                self.advance();
637                self.advance();
638                Token::NotEqual
639            }
640            Some('|') if self.peek(1) == Some('|') => {
641                self.advance();
642                self.advance();
643                Token::Concat
644            }
645            Some('"') => {
646                let ident_val = self.read_string();
647                Token::QuotedIdentifier(ident_val)
648            }
649            Some('$') => {
650                if self.peek_string(6) == "$JSON$" {
651                    let json_content = self.read_json_block();
652                    Token::JsonBlock(json_content)
653                } else {
654                    let ident = self.read_identifier();
655                    Token::Identifier(ident)
656                }
657            }
658            Some('\'') => {
659                let string_val = self.read_string();
660                Token::StringLiteral(string_val)
661            }
662            Some('-') if self.peek(1).is_some_and(char::is_numeric) => {
663                self.advance();
664                let num = self.read_number();
665                Token::NumberLiteral(format!("-{num}"))
666            }
667            Some('-') => {
668                self.advance();
669                Token::Minus
670            }
671            Some(ch) if ch.is_numeric() => {
672                let num = self.read_number();
673                Token::NumberLiteral(num)
674            }
675            Some('#') => {
676                self.advance();
677                let table_name = self.read_identifier();
678                if table_name.is_empty() {
679                    Token::Identifier("#".to_string())
680                } else {
681                    Token::Identifier(format!("#{}", table_name))
682                }
683            }
684            Some(ch) if ch.is_alphabetic() || ch == '_' => {
685                let ident = self.read_identifier();
686                // Handle multi-word keywords like GROUP BY and ORDER BY
687                match ident.to_uppercase().as_str() {
688                    "ORDER" if self.peek_keyword("BY") => {
689                        self.skip_whitespace();
690                        self.read_identifier(); // consume "BY"
691                        Token::OrderBy
692                    }
693                    "GROUP" if self.peek_keyword("BY") => {
694                        self.skip_whitespace();
695                        self.read_identifier(); // consume "BY"
696                        Token::GroupBy
697                    }
698                    _ => Token::from_keyword(&ident).unwrap_or_else(|| Token::Identifier(ident)),
699                }
700            }
701            Some(ch) => {
702                self.advance();
703                Token::Identifier(ch.to_string())
704            }
705        }
706    }
707
708    /// Get next token - dispatches based on lexer mode
709    pub fn next_token(&mut self) -> Token {
710        match self.mode {
711            LexerMode::SkipComments => self.next_token_skip_comments(),
712            LexerMode::PreserveComments => self.next_token_with_comments(),
713        }
714    }
715
716    /// Get next token skipping comments (original behavior)
717    fn next_token_skip_comments(&mut self) -> Token {
718        self.skip_whitespace_and_comments();
719
720        match self.current_char {
721            None => Token::Eof,
722            Some('*') => {
723                self.advance();
724                // Context-sensitive: could be SELECT * or multiplication
725                // The parser will distinguish based on context
726                Token::Star // We'll handle multiplication in parser
727            }
728            Some('+') => {
729                self.advance();
730                Token::Plus
731            }
732            Some('/') => {
733                // Check if this is a comment start
734                if self.peek(1) == Some('*') {
735                    // This shouldn't happen as comments are skipped above,
736                    // but handle it just in case
737                    self.skip_whitespace_and_comments();
738                    return self.next_token();
739                }
740                self.advance();
741                Token::Divide
742            }
743            Some('%') => {
744                self.advance();
745                Token::Modulo
746            }
747            Some('.') => {
748                self.advance();
749                Token::Dot
750            }
751            Some(',') => {
752                self.advance();
753                Token::Comma
754            }
755            Some(';') => {
756                self.advance();
757                Token::Semicolon
758            }
759            Some(':') => {
760                self.advance();
761                Token::Colon
762            }
763            Some('(') => {
764                self.advance();
765                Token::LeftParen
766            }
767            Some(')') => {
768                self.advance();
769                Token::RightParen
770            }
771            Some('=') => {
772                self.advance();
773                Token::Equal
774            }
775            Some('<') => {
776                self.advance();
777                if self.current_char == Some('=') {
778                    self.advance();
779                    Token::LessThanOrEqual
780                } else if self.current_char == Some('>') {
781                    self.advance();
782                    Token::NotEqual
783                } else {
784                    Token::LessThan
785                }
786            }
787            Some('>') => {
788                self.advance();
789                if self.current_char == Some('=') {
790                    self.advance();
791                    Token::GreaterThanOrEqual
792                } else {
793                    Token::GreaterThan
794                }
795            }
796            Some('!') if self.peek(1) == Some('=') => {
797                self.advance();
798                self.advance();
799                Token::NotEqual
800            }
801            Some('|') if self.peek(1) == Some('|') => {
802                self.advance();
803                self.advance();
804                Token::Concat
805            }
806            Some('"') => {
807                // Double quotes = identifier
808                let ident_val = self.read_string();
809                Token::QuotedIdentifier(ident_val)
810            }
811            Some('$') => {
812                // Check if this is $JSON$ delimiter
813                if self.peek_string(6) == "$JSON$" {
814                    let json_content = self.read_json_block();
815                    Token::JsonBlock(json_content)
816                } else {
817                    // Not a JSON block, could be part of identifier or parameter
818                    // For now, treat as identifier start
819                    let ident = self.read_identifier();
820                    Token::Identifier(ident)
821                }
822            }
823            Some('\'') => {
824                // Single quotes = string literal
825                let string_val = self.read_string();
826                Token::StringLiteral(string_val)
827            }
828            Some('-') if self.peek(1) == Some('-') => {
829                // This is a comment, skip it and get next token
830                self.skip_whitespace_and_comments();
831                self.next_token()
832            }
833            Some('-') if self.peek(1).is_some_and(char::is_numeric) => {
834                // Handle negative numbers
835                self.advance(); // skip '-'
836                let num = self.read_number();
837                Token::NumberLiteral(format!("-{num}"))
838            }
839            Some('-') => {
840                // Handle subtraction operator
841                self.advance();
842                Token::Minus
843            }
844            Some(ch) if ch.is_numeric() => {
845                let num = self.read_number();
846                Token::NumberLiteral(num)
847            }
848            Some('#') => {
849                // Temporary table identifier: #tablename
850                self.advance(); // consume #
851                let table_name = self.read_identifier();
852                if table_name.is_empty() {
853                    // Just # by itself
854                    Token::Identifier("#".to_string())
855                } else {
856                    // #tablename
857                    Token::Identifier(format!("#{}", table_name))
858                }
859            }
860            Some(ch) if ch.is_alphabetic() || ch == '_' => {
861                let ident = self.read_identifier();
862                match ident.to_uppercase().as_str() {
863                    "SELECT" => Token::Select,
864                    "FROM" => Token::From,
865                    "WHERE" => Token::Where,
866                    "WITH" => Token::With,
867                    "AND" => Token::And,
868                    "OR" => Token::Or,
869                    "IN" => Token::In,
870                    "NOT" => Token::Not,
871                    "BETWEEN" => Token::Between,
872                    "LIKE" => Token::Like,
873                    "ILIKE" => Token::ILike,
874                    "IS" => Token::Is,
875                    "NULL" => Token::Null,
876                    "ORDER" if self.peek_keyword("BY") => {
877                        self.skip_whitespace();
878                        self.read_identifier(); // consume "BY"
879                        Token::OrderBy
880                    }
881                    "GROUP" if self.peek_keyword("BY") => {
882                        self.skip_whitespace();
883                        self.read_identifier(); // consume "BY"
884                        Token::GroupBy
885                    }
886                    "HAVING" => Token::Having,
887                    "QUALIFY" => Token::Qualify,
888                    "AS" => Token::As,
889                    "ASC" => Token::Asc,
890                    "DESC" => Token::Desc,
891                    "LIMIT" => Token::Limit,
892                    "OFFSET" => Token::Offset,
893                    "INTO" => Token::Into,
894                    "DATETIME" => Token::DateTime,
895                    "CASE" => Token::Case,
896                    "WHEN" => Token::When,
897                    "THEN" => Token::Then,
898                    "ELSE" => Token::Else,
899                    "END" => Token::End,
900                    "DISTINCT" => Token::Distinct,
901                    "EXCLUDE" => Token::Exclude,
902                    "PIVOT" => Token::Pivot,
903                    "UNPIVOT" => Token::Unpivot,
904                    "FOR" => Token::For,
905                    "OVER" => Token::Over,
906                    "PARTITION" => Token::Partition,
907                    "BY" => Token::By,
908                    // Window frame keywords
909                    "ROWS" => Token::Rows,
910                    // Note: RANGE is context-sensitive - it's both a window frame keyword and a table function
911                    // We'll handle this in the parser based on context
912                    "UNBOUNDED" => Token::Unbounded,
913                    "PRECEDING" => Token::Preceding,
914                    "FOLLOWING" => Token::Following,
915                    "CURRENT" => Token::Current,
916                    "ROW" => Token::Row,
917                    // Set operation keywords
918                    "UNION" => Token::Union,
919                    "INTERSECT" => Token::Intersect,
920                    "EXCEPT" => Token::Except,
921                    // Special CTE keywords
922                    "WEB" => Token::Web,
923                    "FILE" => Token::File,
924                    // Row expansion functions
925                    "UNNEST" => Token::Unnest,
926                    // JOIN keywords
927                    "JOIN" => Token::Join,
928                    "INNER" => Token::Inner,
929                    "LEFT" => Token::Left,
930                    "RIGHT" => Token::Right,
931                    "FULL" => Token::Full,
932                    "OUTER" => Token::Outer,
933                    "ON" => Token::On,
934                    "CROSS" => Token::Cross,
935                    _ => Token::Identifier(ident),
936                }
937            }
938            Some(ch) => {
939                self.advance();
940                Token::Identifier(ch.to_string())
941            }
942        }
943    }
944
945    fn peek_keyword(&mut self, keyword: &str) -> bool {
946        let saved_pos = self.position;
947        let saved_char = self.current_char;
948
949        self.skip_whitespace_and_comments();
950        let next_word = self.read_identifier();
951        let matches = next_word.to_uppercase() == keyword;
952
953        // Restore position
954        self.position = saved_pos;
955        self.current_char = saved_char;
956
957        matches
958    }
959
960    #[must_use]
961    pub fn get_position(&self) -> usize {
962        self.position
963    }
964
965    pub fn tokenize_all(&mut self) -> Vec<Token> {
966        let mut tokens = Vec::new();
967        loop {
968            let token = self.next_token();
969            if matches!(token, Token::Eof) {
970                tokens.push(token);
971                break;
972            }
973            tokens.push(token);
974        }
975        tokens
976    }
977
978    pub fn tokenize_all_with_positions(&mut self) -> Vec<(usize, usize, Token)> {
979        let mut tokens = Vec::new();
980        loop {
981            self.skip_whitespace_and_comments();
982            let start_pos = self.position;
983            let token = self.next_token();
984            let end_pos = self.position;
985
986            if matches!(token, Token::Eof) {
987                break;
988            }
989            tokens.push((start_pos, end_pos, token));
990        }
991        tokens
992    }
993
994    /// Tokenize all tokens including comments
995    /// This is useful for formatting tools that need to preserve comments
996    pub fn tokenize_all_with_comments(&mut self) -> Vec<Token> {
997        let mut tokens = Vec::new();
998        loop {
999            let token = self.next_token_with_comments();
1000            if matches!(token, Token::Eof) {
1001                tokens.push(token);
1002                break;
1003            }
1004            tokens.push(token);
1005        }
1006        tokens
1007    }
1008}
1009
1010#[cfg(test)]
1011mod tests {
1012    use super::*;
1013
1014    #[test]
1015    fn test_line_comment_tokenization() {
1016        let sql = "SELECT col1, -- this is a comment\ncol2 FROM table";
1017        let mut lexer = Lexer::new(sql);
1018        let tokens = lexer.tokenize_all_with_comments();
1019
1020        // Find the comment token
1021        let comment_token = tokens.iter().find(|t| matches!(t, Token::LineComment(_)));
1022        assert!(comment_token.is_some(), "Should find line comment token");
1023
1024        if let Some(Token::LineComment(text)) = comment_token {
1025            assert_eq!(text.trim(), "this is a comment");
1026        }
1027    }
1028
1029    #[test]
1030    fn test_block_comment_tokenization() {
1031        let sql = "SELECT /* block comment */ col1 FROM table";
1032        let mut lexer = Lexer::new(sql);
1033        let tokens = lexer.tokenize_all_with_comments();
1034
1035        // Find the comment token
1036        let comment_token = tokens.iter().find(|t| matches!(t, Token::BlockComment(_)));
1037        assert!(comment_token.is_some(), "Should find block comment token");
1038
1039        if let Some(Token::BlockComment(text)) = comment_token {
1040            assert_eq!(text.trim(), "block comment");
1041        }
1042    }
1043
1044    #[test]
1045    fn test_multiple_comments() {
1046        let sql = "-- First comment\nSELECT col1, /* inline */ col2\n-- Second comment\nFROM table";
1047        let mut lexer = Lexer::new(sql);
1048        let tokens = lexer.tokenize_all_with_comments();
1049
1050        let line_comments: Vec<_> = tokens
1051            .iter()
1052            .filter(|t| matches!(t, Token::LineComment(_)))
1053            .collect();
1054        let block_comments: Vec<_> = tokens
1055            .iter()
1056            .filter(|t| matches!(t, Token::BlockComment(_)))
1057            .collect();
1058
1059        assert_eq!(line_comments.len(), 2, "Should find 2 line comments");
1060        assert_eq!(block_comments.len(), 1, "Should find 1 block comment");
1061    }
1062
1063    #[test]
1064    fn test_backwards_compatibility() {
1065        // Test that next_token() still skips comments
1066        let sql = "SELECT -- comment\ncol1 FROM table";
1067        let mut lexer = Lexer::new(sql);
1068        let tokens = lexer.tokenize_all();
1069
1070        // Should NOT contain any comment tokens
1071        let has_comments = tokens
1072            .iter()
1073            .any(|t| matches!(t, Token::LineComment(_) | Token::BlockComment(_)));
1074        assert!(
1075            !has_comments,
1076            "next_token() should skip comments for backwards compatibility"
1077        );
1078
1079        // Should still parse correctly
1080        assert!(tokens.iter().any(|t| matches!(t, Token::Select)));
1081        assert!(tokens.iter().any(|t| matches!(t, Token::From)));
1082    }
1083
1084    // ===== Dual-Mode Lexer Tests (Phase 1) =====
1085
1086    #[test]
1087    fn test_lexer_mode_skip_comments() {
1088        let sql = "SELECT id -- comment\nFROM table";
1089
1090        // SkipComments mode (default)
1091        let mut lexer = Lexer::with_mode(sql, LexerMode::SkipComments);
1092
1093        assert_eq!(lexer.next_token(), Token::Select);
1094        assert_eq!(lexer.next_token(), Token::Identifier("id".into()));
1095        // Comment should be skipped
1096        assert_eq!(lexer.next_token(), Token::From);
1097        assert_eq!(lexer.next_token(), Token::Identifier("table".into()));
1098        assert_eq!(lexer.next_token(), Token::Eof);
1099    }
1100
1101    #[test]
1102    fn test_lexer_mode_preserve_comments() {
1103        let sql = "SELECT id -- comment\nFROM table";
1104
1105        // PreserveComments mode
1106        let mut lexer = Lexer::with_mode(sql, LexerMode::PreserveComments);
1107
1108        assert_eq!(lexer.next_token(), Token::Select);
1109        assert_eq!(lexer.next_token(), Token::Identifier("id".into()));
1110
1111        // Comment should be preserved as a token
1112        let comment_tok = lexer.next_token();
1113        assert!(matches!(comment_tok, Token::LineComment(_)));
1114        if let Token::LineComment(text) = comment_tok {
1115            assert_eq!(text.trim(), "comment");
1116        }
1117
1118        assert_eq!(lexer.next_token(), Token::From);
1119        assert_eq!(lexer.next_token(), Token::Identifier("table".into()));
1120        assert_eq!(lexer.next_token(), Token::Eof);
1121    }
1122
1123    #[test]
1124    fn test_lexer_mode_default_is_skip() {
1125        let sql = "SELECT id -- comment\nFROM table";
1126
1127        // Default (using new()) should skip comments
1128        let mut lexer = Lexer::new(sql);
1129
1130        let mut tok_count = 0;
1131        loop {
1132            let tok = lexer.next_token();
1133            if matches!(tok, Token::Eof) {
1134                break;
1135            }
1136            // Should never see a comment token
1137            assert!(!matches!(
1138                tok,
1139                Token::LineComment(_) | Token::BlockComment(_)
1140            ));
1141            tok_count += 1;
1142        }
1143
1144        // SELECT, id, FROM, table = 4 tokens (no comment)
1145        assert_eq!(tok_count, 4);
1146    }
1147
1148    #[test]
1149    fn test_lexer_mode_block_comments() {
1150        let sql = "SELECT /* block */ id FROM table";
1151
1152        // Skip mode
1153        let mut lexer_skip = Lexer::with_mode(sql, LexerMode::SkipComments);
1154        assert_eq!(lexer_skip.next_token(), Token::Select);
1155        assert_eq!(lexer_skip.next_token(), Token::Identifier("id".into()));
1156        assert_eq!(lexer_skip.next_token(), Token::From);
1157
1158        // Preserve mode
1159        let mut lexer_preserve = Lexer::with_mode(sql, LexerMode::PreserveComments);
1160        assert_eq!(lexer_preserve.next_token(), Token::Select);
1161
1162        let comment_tok = lexer_preserve.next_token();
1163        assert!(matches!(comment_tok, Token::BlockComment(_)));
1164        if let Token::BlockComment(text) = comment_tok {
1165            assert_eq!(text.trim(), "block");
1166        }
1167
1168        assert_eq!(lexer_preserve.next_token(), Token::Identifier("id".into()));
1169    }
1170
1171    #[test]
1172    fn test_lexer_mode_mixed_comments() {
1173        let sql = "-- leading\nSELECT /* inline */ id -- trailing\nFROM table";
1174
1175        let mut lexer = Lexer::with_mode(sql, LexerMode::PreserveComments);
1176
1177        // leading comment
1178        assert!(matches!(lexer.next_token(), Token::LineComment(_)));
1179
1180        // SELECT
1181        assert_eq!(lexer.next_token(), Token::Select);
1182
1183        // inline block comment
1184        assert!(matches!(lexer.next_token(), Token::BlockComment(_)));
1185
1186        // id
1187        assert_eq!(lexer.next_token(), Token::Identifier("id".into()));
1188
1189        // trailing comment
1190        assert!(matches!(lexer.next_token(), Token::LineComment(_)));
1191
1192        // FROM table
1193        assert_eq!(lexer.next_token(), Token::From);
1194        assert_eq!(lexer.next_token(), Token::Identifier("table".into()));
1195        assert_eq!(lexer.next_token(), Token::Eof);
1196    }
1197
1198    #[test]
1199    fn test_pivot_keywords() {
1200        let sql = "PIVOT (MAX(amount) FOR month IN (val1, val2)) UNPIVOT";
1201        let mut lexer = Lexer::new(sql);
1202
1203        // Test individual token recognition
1204        assert_eq!(
1205            lexer.next_token(),
1206            Token::Pivot,
1207            "First token should be PIVOT"
1208        );
1209        assert_eq!(lexer.next_token(), Token::LeftParen);
1210        assert!(matches!(lexer.next_token(), Token::Identifier(_))); // MAX
1211        assert_eq!(lexer.next_token(), Token::LeftParen);
1212        assert!(matches!(lexer.next_token(), Token::Identifier(_))); // amount
1213        assert_eq!(lexer.next_token(), Token::RightParen);
1214        assert_eq!(lexer.next_token(), Token::For, "Should tokenize FOR");
1215        assert!(matches!(lexer.next_token(), Token::Identifier(_))); // month
1216        assert_eq!(lexer.next_token(), Token::In, "Should tokenize IN");
1217        assert_eq!(lexer.next_token(), Token::LeftParen);
1218        assert!(matches!(lexer.next_token(), Token::Identifier(_))); // val1
1219        assert_eq!(lexer.next_token(), Token::Comma);
1220        assert!(matches!(lexer.next_token(), Token::Identifier(_))); // val2
1221        assert_eq!(lexer.next_token(), Token::RightParen);
1222        assert_eq!(lexer.next_token(), Token::RightParen);
1223        assert_eq!(
1224            lexer.next_token(),
1225            Token::Unpivot,
1226            "Should tokenize UNPIVOT"
1227        );
1228        assert_eq!(lexer.next_token(), Token::Eof);
1229    }
1230}
sql_cli/sql/parser/lexer.rs

sql_cli/sql/parser/
lexer.rs
