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
13pub struct TsqlGrammar;
15
16const 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
32static 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 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 self.dispatch_ansi_statement(ctx)
76 }
77 }
78
79 fn parse_table_hint(&self, ctx: &mut ParseContext) -> Option<Segment> {
81 if !ctx.peek_keyword("WITH") {
83 return None;
84 }
85 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 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 ctx.restore(save);
127 return None;
128 }
129
130 Some(Segment::Node(NodeSegment::new(
131 SegmentType::TableHint,
132 children,
133 )))
134 }
135
136 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
199impl TsqlGrammar {
202 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 self.parse_declare_variable(ctx, &mut children);
211
212 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 fn parse_declare_variable(&self, ctx: &mut ParseContext, children: &mut Vec<Segment>) {
235 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 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 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 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 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 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 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 if let Some(dt) = self.parse_data_type(ctx) {
309 children.push(dt);
310 children.extend(eat_trivia_segments(ctx));
311 }
312
313 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 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 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 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 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 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 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 if let Some(cond) = self.parse_expression(ctx) {
390 children.push(cond);
391 }
392 children.extend(eat_trivia_segments(ctx));
393
394 if let Some(stmt) = self.parse_statement(ctx) {
396 children.push(stmt);
397 }
398
399 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 fn parse_begin_block(&self, ctx: &mut ParseContext) -> Option<Segment> {
419 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 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 fn parse_try_catch_block(&self, ctx: &mut ParseContext) -> Option<Segment> {
444 let mut children = Vec::new();
445
446 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 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 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 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 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 if let Some(cond) = self.parse_expression(ctx) {
506 children.push(cond);
507 }
508 children.extend(eat_trivia_segments(ctx));
509
510 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 fn parse_exec_statement(&self, ctx: &mut ParseContext) -> Option<Segment> {
523 let mut children = Vec::new();
524 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 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 if let Some(name) = self.parse_qualified_name(ctx) {
552 children.push(name);
553 }
554 children.extend(eat_trivia_segments(ctx));
555
556 self.parse_exec_params(ctx, &mut children);
558
559 Some(Segment::Node(NodeSegment::new(
560 SegmentType::ExecStatement,
561 children,
562 )))
563 }
564
565 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 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 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 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 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 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 children.extend(trivia);
642 if let Some(expr) = self.parse_expression(ctx) {
643 children.push(expr);
644 }
645 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 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 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 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 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 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 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}