feat(cli): support comment blocks

cli/src/session.rs

@@ -205,39 +205,25 @@ impl Session {
 
         let mut queries = Vec::new();
         let mut tokenizer = Tokenizer::new(line);
-        let mut start = 0;
         let mut in_comment = false;
+        let mut start = 0;
+        let mut comment_block_start = 0;
+
         while let Some(Ok(token)) = tokenizer.next() {
-            match token.kind {
-                TokenKind::CommentBlockStart => {
-                    self.in_comment_block = true;
-                    start = token.span.end;
-                    continue;
-                }
-                TokenKind::CommentBlockEnd => {
-                    if !self.in_comment_block {
-                        panic!("Unexpected comment block end");
-                    }
-                    self.in_comment_block = false;
-                    start = token.span.end;
-                    continue;
-                }
-                _ => {
-                    if self.in_comment_block {
-                        start = token.span.end;
-                        continue;
-                    }
-                }
-            }
             match token.kind {
                 TokenKind::SemiColon => {
-                    let mut sql = self.query.trim().to_owned();
-                    if sql.is_empty() {
+                    if in_comment || self.in_comment_block {
                         continue;
+                    } else {
+                        let mut sql = self.query.trim().to_owned();
+                        if sql.is_empty() {
+                            continue;
+                        }
+                        sql.push(';');
+
+                        queries.push(sql);
+                        self.query.clear();
                     }
-                    sql.push(';');
-                    queries.push(sql);
-                    self.query.clear();
                 }
                 TokenKind::Comment => {
                     in_comment = true;
@@ -249,8 +235,16 @@ impl Session {
                     in_comment = false;
                     self.query.push(' ');
                 }
-                TokenKind::CommentBlockStart | TokenKind::CommentBlockEnd => {
-                    unreachable!("Comment block should be handled before")
+                TokenKind::CommentBlockStart => {
+                    if !self.in_comment_block {
+                        comment_block_start = token.span.start;
+                    }
+                    self.in_comment_block = true;
+                }
+                TokenKind::CommentBlockEnd => {
+                    self.in_comment_block = false;
+                    self.query
+                        .push_str(&line[comment_block_start..token.span.end]);
                 }
                 _ => {
                     if !in_comment && !self.in_comment_block {
@@ -260,6 +254,10 @@ impl Session {
             }
             start = token.span.end;
         }
+
+        if self.in_comment_block {
+            self.query.push_str(&line[comment_block_start..]);
+        }
         queries
     }
 

cli/tests/00-base.result

@@ -6,4 +6,5 @@ a	1	true
 3
 with comment
 3.00	3.00
+Asia/Shanghai
 bye

cli/tests/00-base.sql

@@ -1,3 +1,4 @@
+drop table if exists test;
 create table test(a string, b int, d boolean);
 insert into test values('a', 1, true);
 insert into test values('b', 2, false);
@@ -19,5 +20,7 @@ select 'in comment block';
 
 select 1.00 + 2.00, 3.00;
 
+select/*+ SET_VAR(timezone='Asia/Shanghai') */ timezone();
+
 select 'bye';
 drop table test;