Skip to main content

rigsql_parser/grammar/
tsql.rs

1use std::sync::LazyLock;
2
3use rigsql_core::{NodeSegment, Segment, SegmentType, TokenKind};
4
5use crate::context::ParseContext;
6
7use super::ansi::ANSI_STATEMENT_KEYWORDS;
8use super::{
9    any_token_segment, eat_trivia_segments, parse_comma_separated, parse_statement_list,
10    token_segment, Grammar,
11};
12
13/// TSQL grammar — extends ANSI with SQL Server–specific statements.
14pub struct TsqlGrammar;
15
16/// Keywords that can start a TSQL-specific statement (not in ANSI).
17const TSQL_EXTRA_KEYWORDS: &[&str] = &[
18    "BEGIN",
19    "DECLARE",
20    "EXEC",
21    "EXECUTE",
22    "GO",
23    "IF",
24    "PRINT",
25    "RAISERROR",
26    "RETURN",
27    "SET",
28    "THROW",
29    "WHILE",
30];
31
32/// Full TSQL statement keywords = ANSI + TSQL extras, sorted.
33static TSQL_STATEMENT_KEYWORDS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
34    let mut kws: Vec<&str> = ANSI_STATEMENT_KEYWORDS
35        .iter()
36        .chain(TSQL_EXTRA_KEYWORDS.iter())
37        .copied()
38        .collect();
39    kws.sort_unstable();
40    kws.dedup();
41    kws
42});
43
44impl Grammar for TsqlGrammar {
45    fn statement_keywords(&self) -> &[&str] {
46        &TSQL_STATEMENT_KEYWORDS
47    }
48
49    fn dispatch_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
50        // TSQL-specific statements
51        if ctx.peek_keyword("DECLARE") {
52            self.parse_declare_statement(ctx)
53        } else if ctx.peek_keyword("SET") {
54            self.parse_set_variable_statement(ctx)
55        } else if ctx.peek_keyword("IF") {
56            self.parse_if_statement(ctx)
57        } else if ctx.peek_keyword("BEGIN") {
58            self.parse_begin_block(ctx)
59        } else if ctx.peek_keyword("WHILE") {
60            self.parse_while_statement(ctx)
61        } else if ctx.peek_keyword("EXEC") || ctx.peek_keyword("EXECUTE") {
62            self.parse_exec_statement(ctx)
63        } else if ctx.peek_keyword("RETURN") {
64            self.parse_return_statement(ctx)
65        } else if ctx.peek_keyword("PRINT") {
66            self.parse_print_statement(ctx)
67        } else if ctx.peek_keyword("THROW") {
68            self.parse_throw_statement(ctx)
69        } else if ctx.peek_keyword("RAISERROR") {
70            self.parse_raiserror_statement(ctx)
71        } else if ctx.peek_keyword("GO") {
72            self.parse_go_statement(ctx)
73        } else {
74            // Fall back to ANSI dispatch
75            self.dispatch_ansi_statement(ctx)
76        }
77    }
78
79    /// TSQL table hints: WITH(NOLOCK), WITH(READUNCOMMITTED), etc.
80    fn parse_table_hint(&self, ctx: &mut ParseContext) -> Option<Segment> {
81        // Must start with WITH followed immediately (after optional whitespace) by '('
82        if !ctx.peek_keyword("WITH") {
83            return None;
84        }
85        // Speculatively try: WITH ( hint_list )
86        let save = ctx.save();
87        let with_kw = ctx.eat_keyword("WITH")?;
88        let trivia_after_with = eat_trivia_segments(ctx);
89        if ctx.peek_kind() != Some(TokenKind::LParen) {
90            ctx.restore(save);
91            return None;
92        }
93        let lparen = ctx.advance().unwrap();
94
95        let mut children = Vec::new();
96        children.push(token_segment(with_kw, SegmentType::Keyword));
97        children.extend(trivia_after_with);
98        children.push(token_segment(lparen, SegmentType::LParen));
99        children.extend(eat_trivia_segments(ctx));
100
101        // Parse comma-separated hint keywords
102        let mut first = true;
103        while !ctx.at_eof() && ctx.peek_kind() != Some(TokenKind::RParen) {
104            if !first {
105                if let Some(comma) = ctx.eat_kind(TokenKind::Comma) {
106                    children.push(token_segment(comma, SegmentType::Comma));
107                    children.extend(eat_trivia_segments(ctx));
108                } else {
109                    break;
110                }
111            }
112            first = false;
113            if ctx.peek_kind() == Some(TokenKind::Word) {
114                let hint = ctx.advance().unwrap();
115                children.push(token_segment(hint, SegmentType::Keyword));
116                children.extend(eat_trivia_segments(ctx));
117            } else {
118                break;
119            }
120        }
121
122        if let Some(rparen) = ctx.eat_kind(TokenKind::RParen) {
123            children.push(token_segment(rparen, SegmentType::RParen));
124        } else {
125            // Malformed hint — roll back
126            ctx.restore(save);
127            return None;
128        }
129
130        Some(Segment::Node(NodeSegment::new(
131            SegmentType::TableHint,
132            children,
133        )))
134    }
135
136    /// TSQL override: additionally tracks BEGIN/END block depth and
137    /// stops at GO batch separators.
138    fn consume_until_end(&self, ctx: &mut ParseContext, children: &mut Vec<Segment>) {
139        let mut paren_depth = 0u32;
140        let mut begin_depth = 0u32;
141        let mut case_depth = 0u32;
142        while !ctx.at_eof() {
143            match ctx.peek_kind() {
144                Some(TokenKind::Semicolon) if paren_depth == 0 && begin_depth == 0 => break,
145                Some(TokenKind::LParen) => {
146                    paren_depth += 1;
147                    let token = ctx.advance().unwrap();
148                    children.push(any_token_segment(token));
149                }
150                Some(TokenKind::RParen) => {
151                    paren_depth = paren_depth.saturating_sub(1);
152                    let token = ctx.advance().unwrap();
153                    children.push(any_token_segment(token));
154                }
155                _ => {
156                    let t = ctx.peek().unwrap();
157                    if t.kind == TokenKind::Word {
158                        if t.text.eq_ignore_ascii_case("BEGIN") {
159                            begin_depth += 1;
160                            let token = ctx.advance().unwrap();
161                            children.push(any_token_segment(token));
162                            continue;
163                        } else if t.text.eq_ignore_ascii_case("CASE") {
164                            case_depth += 1;
165                            let token = ctx.advance().unwrap();
166                            children.push(any_token_segment(token));
167                            continue;
168                        } else if t.text.eq_ignore_ascii_case("END") {
169                            if case_depth > 0 {
170                                case_depth -= 1;
171                                let token = ctx.advance().unwrap();
172                                children.push(any_token_segment(token));
173                                continue;
174                            }
175                            if begin_depth > 0 {
176                                begin_depth -= 1;
177                                let token = ctx.advance().unwrap();
178                                children.push(any_token_segment(token));
179                                if begin_depth == 0 && paren_depth == 0 {
180                                    break;
181                                }
182                                continue;
183                            }
184                        } else if t.text.eq_ignore_ascii_case("GO")
185                            && paren_depth == 0
186                            && begin_depth == 0
187                        {
188                            break;
189                        }
190                    }
191                    let token = ctx.advance().unwrap();
192                    children.push(any_token_segment(token));
193                }
194            }
195        }
196    }
197}
198
199// ── TSQL-specific parsing methods ────────────────────────────────
200
201impl TsqlGrammar {
202    /// Parse DECLARE statement: `DECLARE @var TYPE [= expr] [, @var2 TYPE [= expr]]`
203    fn parse_declare_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
204        let mut children = Vec::new();
205        let kw = ctx.eat_keyword("DECLARE")?;
206        children.push(token_segment(kw, SegmentType::Keyword));
207        children.extend(eat_trivia_segments(ctx));
208
209        // First variable declaration
210        self.parse_declare_variable(ctx, &mut children);
211
212        // Additional comma-separated declarations
213        loop {
214            let save = ctx.save();
215            let trivia = eat_trivia_segments(ctx);
216            if let Some(comma) = ctx.eat_kind(TokenKind::Comma) {
217                children.extend(trivia);
218                children.push(token_segment(comma, SegmentType::Comma));
219                children.extend(eat_trivia_segments(ctx));
220                self.parse_declare_variable(ctx, &mut children);
221            } else {
222                ctx.restore(save);
223                break;
224            }
225        }
226
227        Some(Segment::Node(NodeSegment::new(
228            SegmentType::DeclareStatement,
229            children,
230        )))
231    }
232
233    /// Parse a single variable declaration: `@var TYPE [= expr]` or `@var AS TYPE`
234    fn parse_declare_variable(&self, ctx: &mut ParseContext, children: &mut Vec<Segment>) {
235        // @variable name or cursor_name
236        if ctx.peek_kind() == Some(TokenKind::AtSign) {
237            let at = ctx.advance().unwrap();
238            children.push(token_segment(at, SegmentType::Identifier));
239            children.extend(eat_trivia_segments(ctx));
240        } else if ctx.peek_kind() == Some(TokenKind::Word) {
241            // Non-@ cursor name: DECLARE cursor_name CURSOR FOR ...
242            let save = ctx.save();
243            let name = ctx.advance().unwrap();
244            let trivia = eat_trivia_segments(ctx);
245            if ctx.peek_keyword("CURSOR") {
246                children.push(token_segment(name, SegmentType::Identifier));
247                children.extend(trivia);
248            } else {
249                ctx.restore(save);
250            }
251        }
252
253        // Optional AS keyword
254        if ctx.peek_keyword("AS") {
255            let as_kw = ctx.advance().unwrap();
256            children.push(token_segment(as_kw, SegmentType::Keyword));
257            children.extend(eat_trivia_segments(ctx));
258        }
259
260        // CURSOR declaration: DECLARE @cur CURSOR [LOCAL|GLOBAL] [FORWARD_ONLY|SCROLL]
261        //   [STATIC|KEYSET|DYNAMIC|FAST_FORWARD] [READ_ONLY|SCROLL_LOCKS|OPTIMISTIC]
262        //   FOR select_statement
263        if ctx.peek_keyword("CURSOR") {
264            let cursor_kw = ctx.advance().unwrap();
265            children.push(token_segment(cursor_kw, SegmentType::Keyword));
266            children.extend(eat_trivia_segments(ctx));
267
268            // Consume cursor options until FOR
269            while !ctx.at_eof() && !ctx.peek_keyword("FOR") {
270                if ctx.peek_kind() == Some(TokenKind::Semicolon) {
271                    break;
272                }
273                if ctx.peek_kind() == Some(TokenKind::Word) {
274                    let opt = ctx.advance().unwrap();
275                    children.push(token_segment(opt, SegmentType::Keyword));
276                    children.extend(eat_trivia_segments(ctx));
277                } else {
278                    break;
279                }
280            }
281
282            // FOR select_statement
283            if ctx.peek_keyword("FOR") {
284                let for_kw = ctx.advance().unwrap();
285                children.push(token_segment(for_kw, SegmentType::Keyword));
286                children.extend(eat_trivia_segments(ctx));
287                if let Some(sel) = self.parse_select_statement(ctx) {
288                    children.push(sel);
289                }
290            }
291            return;
292        }
293
294        // TABLE variable: DECLARE @t TABLE (...)
295        if ctx.peek_keyword("TABLE") {
296            let table_kw = ctx.advance().unwrap();
297            children.push(token_segment(table_kw, SegmentType::Keyword));
298            children.extend(eat_trivia_segments(ctx));
299            if ctx.peek_kind() == Some(TokenKind::LParen) {
300                if let Some(defs) = self.parse_paren_block(ctx) {
301                    children.push(defs);
302                }
303            }
304            return;
305        }
306
307        // Data type
308        if let Some(dt) = self.parse_data_type(ctx) {
309            children.push(dt);
310            children.extend(eat_trivia_segments(ctx));
311        }
312
313        // Optional = expr (default value)
314        if ctx.peek_kind() == Some(TokenKind::Eq) {
315            let eq = ctx.advance().unwrap();
316            children.push(token_segment(eq, SegmentType::ComparisonOperator));
317            children.extend(eat_trivia_segments(ctx));
318            if let Some(expr) = self.parse_expression(ctx) {
319                children.push(expr);
320            }
321        }
322    }
323
324    /// Parse SET @var = expr or SET option ON/OFF
325    fn parse_set_variable_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
326        let save = ctx.save();
327        let mut children = Vec::new();
328        let kw = ctx.eat_keyword("SET")?;
329        children.push(token_segment(kw, SegmentType::Keyword));
330        children.extend(eat_trivia_segments(ctx));
331
332        // @variable or @@variable → SET @var = expr
333        if ctx.peek_kind() == Some(TokenKind::AtSign) {
334            let at = ctx.advance().unwrap();
335            children.push(token_segment(at, SegmentType::Identifier));
336            children.extend(eat_trivia_segments(ctx));
337
338            // = or += or -= etc.
339            if let Some(kind) = ctx.peek_kind() {
340                if matches!(
341                    kind,
342                    TokenKind::Eq
343                        | TokenKind::Plus
344                        | TokenKind::Minus
345                        | TokenKind::Star
346                        | TokenKind::Slash
347                ) {
348                    let op = ctx.advance().unwrap();
349                    children.push(token_segment(op, SegmentType::Operator));
350                    // Handle compound assignment: +=, -=, etc.
351                    if ctx.peek_kind() == Some(TokenKind::Eq) {
352                        let eq = ctx.advance().unwrap();
353                        children.push(token_segment(eq, SegmentType::Operator));
354                    }
355                    children.extend(eat_trivia_segments(ctx));
356                    if let Some(expr) = self.parse_expression(ctx) {
357                        children.push(expr);
358                    }
359                }
360            }
361
362            return Some(Segment::Node(NodeSegment::new(
363                SegmentType::SetVariableStatement,
364                children,
365            )));
366        }
367
368        // SET OPTION ON/OFF (e.g., SET ANSI_NULLS ON, SET NOCOUNT ON)
369        if ctx.peek_kind() == Some(TokenKind::Word) {
370            self.consume_until_statement_end(ctx, &mut children);
371            return Some(Segment::Node(NodeSegment::new(
372                SegmentType::SetVariableStatement,
373                children,
374            )));
375        }
376
377        ctx.restore(save);
378        None
379    }
380
381    /// Parse IF condition statement [ELSE statement]
382    fn parse_if_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
383        let mut children = Vec::new();
384        let kw = ctx.eat_keyword("IF")?;
385        children.push(token_segment(kw, SegmentType::Keyword));
386        children.extend(eat_trivia_segments(ctx));
387
388        // Condition expression
389        if let Some(cond) = self.parse_expression(ctx) {
390            children.push(cond);
391        }
392        children.extend(eat_trivia_segments(ctx));
393
394        // Then-branch: a single statement or BEGIN...END block
395        if let Some(stmt) = self.parse_statement(ctx) {
396            children.push(stmt);
397        }
398
399        // ELSE branch (optional)
400        children.extend(eat_trivia_segments(ctx));
401        if ctx.peek_keyword("ELSE") {
402            let else_kw = ctx.advance().unwrap();
403            children.push(token_segment(else_kw, SegmentType::Keyword));
404            children.extend(eat_trivia_segments(ctx));
405
406            if let Some(stmt) = self.parse_statement(ctx) {
407                children.push(stmt);
408            }
409        }
410
411        Some(Segment::Node(NodeSegment::new(
412            SegmentType::IfStatement,
413            children,
414        )))
415    }
416
417    /// Parse BEGIN...END block or BEGIN TRY...END TRY BEGIN CATCH...END CATCH
418    fn parse_begin_block(&self, ctx: &mut ParseContext) -> Option<Segment> {
419        // Check for BEGIN TRY / BEGIN CATCH
420        if ctx.peek_keywords(&["BEGIN", "TRY"]) {
421            return self.parse_try_catch_block(ctx);
422        }
423
424        let mut children = Vec::new();
425        let begin_kw = ctx.eat_keyword("BEGIN")?;
426        children.push(token_segment(begin_kw, SegmentType::Keyword));
427
428        parse_statement_list(self, ctx, &mut children, |c| c.peek_keyword("END"));
429
430        // END
431        children.extend(eat_trivia_segments(ctx));
432        if let Some(end_kw) = ctx.eat_keyword("END") {
433            children.push(token_segment(end_kw, SegmentType::Keyword));
434        }
435
436        Some(Segment::Node(NodeSegment::new(
437            SegmentType::BeginEndBlock,
438            children,
439        )))
440    }
441
442    /// Parse BEGIN TRY...END TRY BEGIN CATCH...END CATCH
443    fn parse_try_catch_block(&self, ctx: &mut ParseContext) -> Option<Segment> {
444        let mut children = Vec::new();
445
446        // BEGIN TRY
447        let begin_kw = ctx.eat_keyword("BEGIN")?;
448        children.push(token_segment(begin_kw, SegmentType::Keyword));
449        children.extend(eat_trivia_segments(ctx));
450        let try_kw = ctx.eat_keyword("TRY")?;
451        children.push(token_segment(try_kw, SegmentType::Keyword));
452
453        parse_statement_list(self, ctx, &mut children, |c| {
454            c.peek_keywords(&["END", "TRY"])
455        });
456
457        // END TRY
458        children.extend(eat_trivia_segments(ctx));
459        if let Some(end_kw) = ctx.eat_keyword("END") {
460            children.push(token_segment(end_kw, SegmentType::Keyword));
461            children.extend(eat_trivia_segments(ctx));
462        }
463        if let Some(try_kw) = ctx.eat_keyword("TRY") {
464            children.push(token_segment(try_kw, SegmentType::Keyword));
465        }
466
467        // BEGIN CATCH
468        children.extend(eat_trivia_segments(ctx));
469        if let Some(begin_kw) = ctx.eat_keyword("BEGIN") {
470            children.push(token_segment(begin_kw, SegmentType::Keyword));
471            children.extend(eat_trivia_segments(ctx));
472            if let Some(catch_kw) = ctx.eat_keyword("CATCH") {
473                children.push(token_segment(catch_kw, SegmentType::Keyword));
474            }
475
476            parse_statement_list(self, ctx, &mut children, |c| {
477                c.peek_keywords(&["END", "CATCH"])
478            });
479
480            // END CATCH
481            children.extend(eat_trivia_segments(ctx));
482            if let Some(end_kw) = ctx.eat_keyword("END") {
483                children.push(token_segment(end_kw, SegmentType::Keyword));
484                children.extend(eat_trivia_segments(ctx));
485            }
486            if let Some(catch_kw) = ctx.eat_keyword("CATCH") {
487                children.push(token_segment(catch_kw, SegmentType::Keyword));
488            }
489        }
490
491        Some(Segment::Node(NodeSegment::new(
492            SegmentType::TryCatchBlock,
493            children,
494        )))
495    }
496
497    /// Parse WHILE condition statement
498    fn parse_while_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
499        let mut children = Vec::new();
500        let kw = ctx.eat_keyword("WHILE")?;
501        children.push(token_segment(kw, SegmentType::Keyword));
502        children.extend(eat_trivia_segments(ctx));
503
504        // Condition
505        if let Some(cond) = self.parse_expression(ctx) {
506            children.push(cond);
507        }
508        children.extend(eat_trivia_segments(ctx));
509
510        // Body: usually BEGIN...END
511        if let Some(stmt) = self.parse_statement(ctx) {
512            children.push(stmt);
513        }
514
515        Some(Segment::Node(NodeSegment::new(
516            SegmentType::WhileStatement,
517            children,
518        )))
519    }
520
521    /// Parse EXEC/EXECUTE proc_name [params]
522    fn parse_exec_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
523        let mut children = Vec::new();
524        // EXEC or EXECUTE
525        let kw = if ctx.peek_keyword("EXEC") {
526            ctx.eat_keyword("EXEC")
527        } else {
528            ctx.eat_keyword("EXECUTE")
529        };
530        let kw = kw?;
531        children.push(token_segment(kw, SegmentType::Keyword));
532        children.extend(eat_trivia_segments(ctx));
533
534        // Optional @retval =
535        let save = ctx.save();
536        if ctx.peek_kind() == Some(TokenKind::AtSign) {
537            let at = ctx.advance().unwrap();
538            let trivia = eat_trivia_segments(ctx);
539            if ctx.peek_kind() == Some(TokenKind::Eq) {
540                children.push(token_segment(at, SegmentType::Identifier));
541                children.extend(trivia);
542                let eq = ctx.advance().unwrap();
543                children.push(token_segment(eq, SegmentType::Operator));
544                children.extend(eat_trivia_segments(ctx));
545            } else {
546                ctx.restore(save);
547            }
548        }
549
550        // Procedure name (possibly qualified)
551        if let Some(name) = self.parse_qualified_name(ctx) {
552            children.push(name);
553        }
554        children.extend(eat_trivia_segments(ctx));
555
556        // Parameters: comma-separated expressions / @param = expr
557        self.parse_exec_params(ctx, &mut children);
558
559        Some(Segment::Node(NodeSegment::new(
560            SegmentType::ExecStatement,
561            children,
562        )))
563    }
564
565    /// Parse EXEC parameters: comma-separated, optionally @param = expr
566    fn parse_exec_params(&self, ctx: &mut ParseContext, children: &mut Vec<Segment>) {
567        if ctx.at_eof()
568            || ctx.peek_kind() == Some(TokenKind::Semicolon)
569            || self.peek_statement_start(ctx)
570        {
571            return;
572        }
573        parse_comma_separated(ctx, children, |c| self.parse_expression(c));
574    }
575
576    /// Parse RETURN [expr]
577    fn parse_return_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
578        let mut children = Vec::new();
579        let kw = ctx.eat_keyword("RETURN")?;
580        children.push(token_segment(kw, SegmentType::Keyword));
581
582        // Optional return value
583        let save = ctx.save();
584        let trivia = eat_trivia_segments(ctx);
585        if !ctx.at_eof()
586            && ctx.peek_kind() != Some(TokenKind::Semicolon)
587            && !self.peek_statement_start(ctx)
588        {
589            children.extend(trivia);
590            if let Some(expr) = self.parse_expression(ctx) {
591                children.push(expr);
592            }
593        } else {
594            ctx.restore(save);
595        }
596
597        Some(Segment::Node(NodeSegment::new(
598            SegmentType::ReturnStatement,
599            children,
600        )))
601    }
602
603    /// Parse PRINT expr
604    fn parse_print_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
605        let mut children = Vec::new();
606        let kw = ctx.eat_keyword("PRINT")?;
607        children.push(token_segment(kw, SegmentType::Keyword));
608        children.extend(eat_trivia_segments(ctx));
609
610        if let Some(expr) = self.parse_expression(ctx) {
611            children.push(expr);
612        }
613
614        Some(Segment::Node(NodeSegment::new(
615            SegmentType::PrintStatement,
616            children,
617        )))
618    }
619
620    /// Parse THROW [number, message, state]
621    fn parse_throw_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
622        let mut children = Vec::new();
623        let kw = ctx.eat_keyword("THROW")?;
624        children.push(token_segment(kw, SegmentType::Keyword));
625
626        // THROW with no arguments (re-throw in CATCH block)
627        let save = ctx.save();
628        let trivia = eat_trivia_segments(ctx);
629        if ctx.at_eof()
630            || ctx.peek_kind() == Some(TokenKind::Semicolon)
631            || self.peek_statement_start(ctx)
632        {
633            ctx.restore(save);
634            return Some(Segment::Node(NodeSegment::new(
635                SegmentType::ThrowStatement,
636                children,
637            )));
638        }
639
640        // THROW error_number, message, state
641        children.extend(trivia);
642        if let Some(expr) = self.parse_expression(ctx) {
643            children.push(expr);
644        }
645        // Consume remaining comma-separated args
646        for _ in 0..2 {
647            let save2 = ctx.save();
648            let trivia2 = eat_trivia_segments(ctx);
649            if let Some(comma) = ctx.eat_kind(TokenKind::Comma) {
650                children.extend(trivia2);
651                children.push(token_segment(comma, SegmentType::Comma));
652                children.extend(eat_trivia_segments(ctx));
653                if let Some(expr) = self.parse_expression(ctx) {
654                    children.push(expr);
655                }
656            } else {
657                ctx.restore(save2);
658                break;
659            }
660        }
661
662        Some(Segment::Node(NodeSegment::new(
663            SegmentType::ThrowStatement,
664            children,
665        )))
666    }
667
668    /// Parse RAISERROR(msg, severity, state) [WITH option]
669    fn parse_raiserror_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
670        let mut children = Vec::new();
671        let kw = ctx.eat_keyword("RAISERROR")?;
672        children.push(token_segment(kw, SegmentType::Keyword));
673        children.extend(eat_trivia_segments(ctx));
674
675        // Arguments in parens
676        if ctx.peek_kind() == Some(TokenKind::LParen) {
677            if let Some(args) = self.parse_paren_block(ctx) {
678                children.push(args);
679            }
680        }
681
682        // Optional WITH option
683        children.extend(eat_trivia_segments(ctx));
684        if ctx.peek_keyword("WITH") {
685            let with_kw = ctx.advance().unwrap();
686            children.push(token_segment(with_kw, SegmentType::Keyword));
687            children.extend(eat_trivia_segments(ctx));
688            // Options: NOWAIT, LOG, SETERROR
689            while ctx.peek_kind() == Some(TokenKind::Word) {
690                let opt = ctx.advance().unwrap();
691                children.push(token_segment(opt, SegmentType::Keyword));
692                let save = ctx.save();
693                let trivia = eat_trivia_segments(ctx);
694                if ctx.peek_kind() == Some(TokenKind::Comma) {
695                    children.extend(trivia);
696                    let comma = ctx.advance().unwrap();
697                    children.push(token_segment(comma, SegmentType::Comma));
698                    children.extend(eat_trivia_segments(ctx));
699                } else {
700                    ctx.restore(save);
701                    break;
702                }
703            }
704        }
705
706        Some(Segment::Node(NodeSegment::new(
707            SegmentType::RaiserrorStatement,
708            children,
709        )))
710    }
711
712    /// Parse GO batch separator
713    fn parse_go_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
714        let mut children = Vec::new();
715        let kw = ctx.eat_keyword("GO")?;
716        children.push(token_segment(kw, SegmentType::Keyword));
717
718        // Optional count: GO 5
719        let save = ctx.save();
720        let trivia = eat_trivia_segments(ctx);
721        if ctx.peek_kind() == Some(TokenKind::NumberLiteral) {
722            children.extend(trivia);
723            let num = ctx.advance().unwrap();
724            children.push(token_segment(num, SegmentType::NumericLiteral));
725        } else {
726            ctx.restore(save);
727        }
728
729        Some(Segment::Node(NodeSegment::new(
730            SegmentType::GoStatement,
731            children,
732        )))
733    }
734}