support for parameter references

sigmacomputing · Mar 17, 2024 · 1a85f68 · 1a85f68
1 parent 2730712
commit 1a85f68
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 1 deletion.
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -374,6 +374,8 @@ pub enum Expr {
     Identifier(Ident),
     /// Multi-part identifier, e.g. `table_alias.column` or `schema.table.col`
     CompoundIdentifier(Vec<Ident>),
+    /// A reference to a Sigma scalar value, e.g. `@sigma.my_parameter`.
+    SigmaParameter(Ident),
     /// JSON access (postgres)  eg: data->'tags'
     JsonAccess {
         left: Box<Expr>,
@@ -752,6 +754,7 @@ impl fmt::Display for Expr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             Expr::Identifier(s) => write!(f, "{s}"),
+            Expr::SigmaParameter(s) => write!(f, "@sigma.{s}"),
             Expr::MapAccess { column, keys } => {
                 write!(f, "{column}")?;
                 for k in keys {

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -1050,6 +1050,10 @@ impl<'a> Parser<'a> {
             }, // End of Token::Word
             // array `[1, 2, 3]`
             Token::LBracket => self.parse_array_expr(false),
+            Token::AtSign if self.parse_keyword(Keyword::SIGMA) => {
+                self.expect_token(&Token::Period)?;
+                Ok(Expr::SigmaParameter(self.parse_identifier(false)?))
+            }
             tok @ Token::Minus | tok @ Token::Plus => {
                 let op = if tok == Token::Plus {
                     UnaryOperator::Plus

diff --git a/tests/sqlparser_sigma.rs b/tests/sqlparser_sigma.rs
@@ -15,7 +15,7 @@ fn snowflake() -> TestedDialects {
 }
 #[test]
 fn parse_sigma() {
-    let sql = "SELECT my_column FROM @sigma.my_element";
+    let sql = "SELECT my_column FROM @sigma.my_element WHERE my_column <> @sigma.param_filter";
     let select = snowflake().verified_only_select(sql);
     assert_eq!(
         select.from,
@@ -26,5 +26,13 @@ fn parse_sigma() {
             },
             joins: vec![]
         }]
+    );
+    assert_eq!(
+        select.selection,
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("my_column"))),
+            op: BinaryOperator::NotEq,
+            right: Box::new(Expr::SigmaParameter(Ident::new("param_filter"))),
+        })
     )
 }