chore(query): improve variable parser (#16599)

* chore(query): improve variable parser * chore(query): improve variable parser * chore(query): improve variable parser
databendlabs · Oct 13, 2024 · 01bc7ed · 01bc7ed
1 parent 96106c4
commit 01bc7ed
Show file tree

Hide file tree

Showing 6 changed files with 21 additions and 4 deletions.
diff --git a/src/query/ast/src/parser/common.rs b/src/query/ast/src/parser/common.rs
@@ -294,7 +294,7 @@ pub fn column_id(i: Input) -> IResult<ColumnID> {
 }
 
 pub fn variable_ident(i: Input) -> IResult<String> {
-    map(rule! { "$" ~ ^#plain_ident }, |(_, name)| name.name)(i)
+    map(rule! { IdentVariable }, |t| t.text()[1..].to_string())(i)
 }
 
 /// Parse one to two idents separated by a dot, fulfilling from the right.

diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -146,6 +146,9 @@ pub enum TokenKind {
     #[regex(r#"[_a-zA-Z][_$a-zA-Z0-9]*"#)]
     Ident,
 
+    #[regex(r#"\$[_a-zA-Z][_$a-zA-Z0-9]*"#)]
+    IdentVariable,
+
     #[regex(r#"\$[0-9]+"#)]
     ColumnPosition,
 

diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs
@@ -788,7 +788,7 @@ fn test_statement() {
         r#"DROP FUNCTION binary_reverse;"#,
         r#"DROP FUNCTION isnotempty;"#,
         r#"
-            EXECUTE IMMEDIATE 
+            EXECUTE IMMEDIATE
             $$
             BEGIN
                 LOOP
@@ -1239,6 +1239,7 @@ fn test_expr_error() {
         r#"1 a"#,
         r#"CAST(col1)"#,
         r#"a.add(b)"#,
+        r#"$ abc + 3"#,
         r#"[ x * 100 FOR x in [1,2,3] if x % 2 = 0 ]"#,
         r#"
             G.E.B IS NOT NULL

diff --git a/src/query/ast/tests/it/testdata/expr-error.txt b/src/query/ast/tests/it/testdata/expr-error.txt
@@ -72,6 +72,19 @@ error:
   | while parsing expression
 
 
+---------- Input ----------
+$ abc + 3
+---------- Output ---------
+error: 
+  --> SQL:1:1
+  |
+1 | $ abc + 3
+  | ^
+  | |
+  | unexpected `$`, expecting `IS`, `IN`, `EXISTS`, `BETWEEN`, `+`, `-`, `*`, `/`, `//`, `DIV`, `%`, `||`, `<->`, `>`, `<`, `>=`, `<=`, `=`, `<>`, `!=`, `^`, `AND`, `OR`, `XOR`, `LIKE`, `NOT`, `REGEXP`, `RLIKE`, `SOUNDS`, <BitWiseOr>, <BitWiseAnd>, <BitWiseXor>, <ShiftLeft>, <ShiftRight>, `->`, `->>`, `#>`, `#>>`, `?`, `?|`, `?&`, `@>`, `<@`, `@?`, `@@`, `#-`, <Factorial>, <SquareRoot>, <BitWiseNot>, <CubeRoot>, <Abs>, `CAST`, `TRY_CAST`, `DATE_ADD`, `DATE_DIFF`, `DATE_SUB`, `DATE_TRUNC`, `DATE`, `TIMESTAMP`, `INTERVAL`, or 31 more ...
+  | while parsing expression
+
+
 ---------- Input ----------
 [ x * 100 FOR x in [1,2,3] if x % 2 = 0 ]
 ---------- Output ---------

diff --git a/src/query/ast/tests/it/testdata/stmt-error.txt b/src/query/ast/tests/it/testdata/stmt-error.txt
@@ -892,7 +892,7 @@ error:
   --> SQL:1:23
   |
 1 | drop table IDENTIFIER(a)
-  | ----                  ^ unexpected `a`, expecting `$` or `:`
+  | ----                  ^ unexpected `a`, expecting `IdentVariable` or `:`
   | |                      
   | while parsing `DROP TABLE [IF EXISTS] [<database>.]<table>`
 

diff --git a/src/query/ast/tests/it/testdata/stmt.txt b/src/query/ast/tests/it/testdata/stmt.txt
@@ -22377,7 +22377,7 @@ DropUDF {
 
 
 ---------- Input ----------
-EXECUTE IMMEDIATE 
+EXECUTE IMMEDIATE
 $$
 BEGIN
     LOOP