From 4b9751658c5cceafdab65ca04a47f4bde7990451 Mon Sep 17 00:00:00 2001
From: Yang Xiufeng <yangxiufeng.c@gmail.com>
Date: Wed, 25 Sep 2024 19:24:52 +0800
Subject: [PATCH] feat: list stage name to plain_identifier.

---
 src/query/ast/src/parser/common.rs | 2 +-
 src/query/ast/src/parser/token.rs  | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/query/ast/src/parser/common.rs b/src/query/ast/src/parser/common.rs
index 7b4abc56b4849..6971f0d03c644 100644
--- a/src/query/ast/src/parser/common.rs
+++ b/src/query/ast/src/parser/common.rs
@@ -102,7 +102,7 @@ pub fn stage_name(i: Input) -> IResult<Identifier> {
     });
 
     rule!(
-        #ident
+        #plain_identifier(|token| token.is_reserved_ident(false))(i)
         | #anonymous_stage
     )(i)
 }
diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
index 5b9f26ed6bcbc..0063a3738191f 100644
--- a/src/query/ast/src/parser/token.rs
+++ b/src/query/ast/src/parser/token.rs
@@ -157,7 +157,10 @@ pub enum TokenKind {
     #[regex(r#"\$\$([^\$]|(\$[^\$]))*\$\$"#)]
     LiteralCodeString,
 
-    #[regex(r#"@([^\s`;'"()]|\\\s|\\'|\\"|\\\\)+"#)]
+    #[regex(r#"@[_a-zA-Z][_$a-zA-Z0-9]*"#)]
+    #[regex(r#"@`[^`]*`"#)]
+    #[regex(r#"@"([^"\\]|\\.|"")*""#)]
+    #[regex(r#"@'([^'\\]|\\.|'')*'"#)]
     LiteralAtString,
 
     #[regex(r"[xX]'[a-fA-F0-9]*'")]