feat(query): add identifier function from session variable (databendl…

…abs#16264) * feat(query): add identifier function from session variable * feat(query): add identifier function from session variable * feat(query): add identifier function from session variable * fix parsing hole and improve error report * update * update * update * update * update --------- Co-authored-by: andylokandy <[email protected]>
ClSlaid · Aug 17, 2024 · a0a8b31 · a0a8b31
1 parent 695e169
commit a0a8b31
Show file tree

Hide file tree

Showing 22 changed files with 1,359 additions and 1,186 deletions.
diff --git a/src/meta/api/src/txn_backoff.rs b/src/meta/api/src/txn_backoff.rs
@@ -155,7 +155,7 @@ mod tests {
         let elapsed = now.elapsed().as_secs_f64();
         println!("elapsed: {elapsed}");
         assert!(
-            (0.041..0.080).contains(&elapsed),
+            (0.041..0.090).contains(&elapsed),
             "{} is expected to be 2 + 5 + 10 + 14 + 20 milliseconds",
             elapsed
         );

diff --git a/src/query/ast/src/ast/common.rs b/src/query/ast/src/ast/common.rs
@@ -29,20 +29,37 @@ pub struct Identifier {
     pub span: Span,
     pub name: String,
     pub quote: Option<char>,
-    pub is_hole: bool,
+    #[drive(skip)]
+    pub ident_type: IdentifierType,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
+pub enum IdentifierType {
+    #[default]
+    None,
+    Hole,
+    Variable,
 }
 
 impl Identifier {
     pub fn is_quoted(&self) -> bool {
         self.quote.is_some()
     }
 
+    pub fn is_hole(&self) -> bool {
+        self.ident_type == IdentifierType::Hole
+    }
+
+    pub fn is_variable(&self) -> bool {
+        self.ident_type == IdentifierType::Variable
+    }
+
     pub fn from_name(span: Span, name: impl Into<String>) -> Self {
         Self {
             span,
             name: name.into(),
             quote: None,
-            is_hole: false,
+            ident_type: IdentifierType::None,
         }
     }
 
@@ -51,15 +68,17 @@ impl Identifier {
             span,
             name: name.into(),
             quote,
-            is_hole: false,
+            ident_type: IdentifierType::None,
         }
     }
 }
 
 impl Display for Identifier {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        if self.is_hole {
+        if self.is_hole() {
             write!(f, "IDENTIFIER(:{})", self.name)
+        } else if self.is_variable() {
+            write!(f, "IDENTIFIER(${})", self.name)
         } else if let Some(quote) = self.quote {
             write!(f, "{}", QuotedIdent(&self.name, quote))
         } else {

diff --git a/src/query/ast/src/parser/common.rs b/src/query/ast/src/parser/common.rs
@@ -27,6 +27,7 @@ use crate::ast::quote::QuotedIdent;
 use crate::ast::ColumnID;
 use crate::ast::DatabaseRef;
 use crate::ast::Identifier;
+use crate::ast::IdentifierType;
 use crate::ast::SetType;
 use crate::ast::TableRef;
 use crate::parser::input::Input;
@@ -127,7 +128,7 @@ fn plain_identifier(
                 span: transform_span(&[token.clone()]),
                 name: token.text().to_string(),
                 quote: None,
-                is_hole: false,
+                ident_type: IdentifierType::None,
             },
         )(i)
     }
@@ -152,7 +153,7 @@ fn quoted_identifier(i: Input) -> IResult<Identifier> {
                 span: transform_span(&[token.clone()]),
                 name: ident,
                 quote: Some(quote),
-                is_hole: false,
+                ident_type: IdentifierType::None,
             }))
         } else {
             Err(nom::Err::Error(Error::from_error_kind(
@@ -166,17 +167,31 @@ fn quoted_identifier(i: Input) -> IResult<Identifier> {
 fn identifier_hole(i: Input) -> IResult<Identifier> {
     check_template_mode(map(
         consumed(rule! {
-            IDENTIFIER ~ ^"(" ~ ^#template_hole ~ ^")"
+            IDENTIFIER ~ ^"(" ~ #template_hole ~ ^")"
         }),
-        |(span, (_, _, (_, name), _))| Identifier {
+        |(span, (_, _, name, _))| Identifier {
             span: transform_span(span.tokens),
             name,
             quote: None,
-            is_hole: true,
+            ident_type: IdentifierType::Hole,
         },
     ))(i)
 }
 
+fn identifier_variable(i: Input) -> IResult<Identifier> {
+    map(
+        consumed(rule! {
+            IDENTIFIER ~ ^"(" ~ ^#variable_ident ~ ^")"
+        }),
+        |(span, (_, _, name, _))| Identifier {
+            span: transform_span(span.tokens),
+            name,
+            quote: None,
+            ident_type: IdentifierType::Variable,
+        },
+    )(i)
+}
+
 fn non_reserved_identifier(
     is_reserved_keyword: fn(&TokenKind) -> bool,
 ) -> impl FnMut(Input) -> IResult<Identifier> {
@@ -185,6 +200,7 @@ fn non_reserved_identifier(
             #plain_identifier(is_reserved_keyword)
             | #quoted_identifier
             | #identifier_hole
+            | #identifier_variable
         )(i)
     }
 }
@@ -277,11 +293,11 @@ pub fn column_id(i: Input) -> IResult<ColumnID> {
     ))(i)
 }
 
-pub fn variable_ident(i: Input) -> IResult<Identifier> {
-    map(rule! { VariableAccess }, |token| {
-        let name = token.text().to_string();
-        Identifier::from_name(Some(token.span), &name[1..])
-    })(i)
+pub fn variable_ident(i: Input) -> IResult<String> {
+    map(
+        rule! { "$" ~ ^#plain_identifier(|token| token.is_reserved_ident(false)) },
+        |(_, name)| name.name,
+    )(i)
 }
 
 /// Parse one to two idents separated by a dot, fulfilling from the right.
@@ -564,12 +580,12 @@ where F: nom::Parser<Input<'a>, O, Error<'a>> {
     }
 }
 
-pub fn template_hole(i: Input) -> IResult<(Span, String)> {
+pub fn template_hole(i: Input) -> IResult<String> {
     check_template_mode(map(
-        consumed(rule! {
+        rule! {
             ":" ~ ^#plain_identifier(|token| token.is_reserved_ident(false))
-        }),
-        |(span, (_, ident))| (transform_span(span.tokens), ident.name),
+        },
+        |(_, name)| name.name,
     ))(i)
 }
 

diff --git a/src/query/ast/src/parser/expr.rs b/src/query/ast/src/parser/expr.rs
@@ -104,7 +104,7 @@ pub fn subexpr(min_precedence: u32) -> impl FnMut(Input) -> IResult<Expr> {
                     ExprElement::MapAccess {
                         accessor: MapAccessor::Colon { key },
                     } => {
-                        if !key.is_quoted() && !key.is_hole {
+                        if !key.is_quoted() && !key.is_hole() {
                             *elem = ExprElement::Hole {
                                 name: key.to_string(),
                             };
@@ -182,7 +182,7 @@ pub enum ExprElement {
     UnaryOp {
         op: UnaryOperator,
     },
-    VariableAccess(Identifier),
+    VariableAccess(String),
     /// `CAST` expression, like `CAST(expr AS target_type)`
     Cast {
         expr: Box<Expr>,
@@ -650,7 +650,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, ExprElement>>> PrattParser<I> for ExprP
                     name: Identifier::from_name(transform_span(elem.span.tokens), "getvariable"),
                     args: vec![Expr::Literal {
                         span: transform_span(elem.span.tokens),
-                        value: Literal::String(name.to_string()),
+                        value: Literal::String(name),
                     }],
                     params: vec![],
                     window: None,

diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -149,9 +149,6 @@ pub enum TokenKind {
     #[regex(r#"\$[0-9]+"#)]
     ColumnPosition,
 
-    #[regex(r#"\$[_a-zA-Z][_$a-zA-Z0-9]*"#)]
-    VariableAccess,
-
     #[regex(r#"`[^`]*`"#)]
     #[regex(r#""([^"\\]|\\.|"")*""#)]
     #[regex(r#"'([^'\\]|\\.|'')*'"#)]
@@ -239,6 +236,8 @@ pub enum TokenKind {
     LBrace,
     #[token("}")]
     RBrace,
+    #[token("$")]
+    Dollar,
     #[token("->")]
     RArrow,
     #[token("->>")]
@@ -1376,6 +1375,7 @@ impl TokenKind {
                 | Factorial
                 | LBrace
                 | RBrace
+                | Dollar
                 | RArrow
                 | LongRArrow
                 | HashRArrow

diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs
@@ -91,7 +91,7 @@ fn run_parser_with_dialect<P, O>(
 #[test]
 fn test_statement() {
     let mut mint = Mint::new("tests/it/testdata");
-    let file = &mut mint.new_goldenfile("statement.txt").unwrap();
+    let file = &mut mint.new_goldenfile("stmt.txt").unwrap();
     let cases = &[
         r#"show databases"#,
         r#"show databases format TabSeparatedWithNamesAndTypes;"#,
@@ -580,6 +580,7 @@ fn test_statement() {
         r#"SHOW GRANTS OF SHARE t;"#,
         r#"UPDATE db1.tb1 set a = a + 1, b = 2 WHERE c > 3;"#,
         r#"select $abc + 3"#,
+        r#"select IDENTIFIER($abc)"#,
         r#"SET max_threads = 10;"#,
         r#"SET max_threads = 10*2;"#,
         r#"SET global (max_threads, max_memory_usage) = (10*2, 10*4);"#,
@@ -864,7 +865,7 @@ fn test_statement() {
 #[test]
 fn test_statement_error() {
     let mut mint = Mint::new("tests/it/testdata");
-    let file = &mut mint.new_goldenfile("statement-error.txt").unwrap();
+    let file = &mut mint.new_goldenfile("stmt-error.txt").unwrap();
 
     let cases = &[
         r#"create table a.b (c integer not null 1, b float(10))"#,