From 27307128fd08cd226a33e1c513bd8bb833575027 Mon Sep 17 00:00:00 2001
From: Joey Hain <hain.joseph@gmail.com>
Date: Sat, 9 Mar 2024 13:30:15 -0800
Subject: [PATCH] sigma element table factor

---
 src/ast/query.rs         | 12 ++++++++++++
 src/keywords.rs          |  1 +
 src/parser/mod.rs        | 13 +++++++++++++
 tests/sqlparser_sigma.rs | 30 ++++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+)
 create mode 100644 tests/sqlparser_sigma.rs

diff --git a/src/ast/query.rs b/src/ast/query.rs
index d7037ce4c..04309d9b7 100644
--- a/src/ast/query.rs
+++ b/src/ast/query.rs
@@ -761,6 +761,11 @@ pub enum TableFactor {
         /// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
         partitions: Vec<Ident>,
     },
+    /// A reference to an element in a Sigma workbook.
+    SigmaElement {
+        element: Ident,
+        alias: Option<TableAlias>,
+    },
     Derived {
         lateral: bool,
         subquery: Box<Query>,
@@ -887,6 +892,13 @@ impl fmt::Display for TableFactor {
                 }
                 Ok(())
             }
+            TableFactor::SigmaElement { element, alias } => {
+                write!(f, "@sigma.{element}")?;
+                if let Some(alias) = alias {
+                    write!(f, " AS {alias}")?;
+                }
+                Ok(())
+            }
             TableFactor::Derived {
                 lateral,
                 subquery,
diff --git a/src/keywords.rs b/src/keywords.rs
index 3621961bc..b3208ffeb 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -612,6 +612,7 @@ define_keywords!(
     SETS,
     SHARE,
     SHOW,
+    SIGMA,
     SIMILAR,
     SKIP,
     SLOW,
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index f9ca62f8d..14205ab87 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -7985,6 +7985,7 @@ impl<'a> Parser<'a> {
                     match &mut table_and_joins.relation {
                         TableFactor::Derived { alias, .. }
                         | TableFactor::Table { alias, .. }
+                        | TableFactor::SigmaElement { alias, .. }
                         | TableFactor::Function { alias, .. }
                         | TableFactor::UNNEST { alias, .. }
                         | TableFactor::JsonTable { alias, .. }
@@ -8062,6 +8063,18 @@ impl<'a> Parser<'a> {
                 columns,
                 alias,
             })
+        } else if self
+            .maybe_parse(|p| {
+                p.expect_token(&Token::AtSign)?;
+                p.expect_keyword(Keyword::SIGMA)?;
+                Ok(())
+            })
+            .is_some()
+        {
+            self.expect_token(&Token::Period)?;
+            let element = self.parse_identifier(true)?;
+            let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
+            Ok(TableFactor::SigmaElement { element, alias })
         } else {
             let name = self.parse_object_name(true)?;
 
diff --git a/tests/sqlparser_sigma.rs b/tests/sqlparser_sigma.rs
new file mode 100644
index 000000000..38dd7060d
--- /dev/null
+++ b/tests/sqlparser_sigma.rs
@@ -0,0 +1,30 @@
+#![warn(clippy::all)]
+
+use sqlparser::ast::*;
+use sqlparser::dialect::SnowflakeDialect;
+use test_utils::*;
+
+#[macro_use]
+mod test_utils;
+
+fn snowflake() -> TestedDialects {
+    TestedDialects {
+        dialects: vec![Box::new(SnowflakeDialect {})],
+        options: None,
+    }
+}
+#[test]
+fn parse_sigma() {
+    let sql = "SELECT my_column FROM @sigma.my_element";
+    let select = snowflake().verified_only_select(sql);
+    assert_eq!(
+        select.from,
+        vec![TableWithJoins {
+            relation: TableFactor::SigmaElement {
+                element: Ident::new("my_element"),
+                alias: None
+            },
+            joins: vec![]
+        }]
+    )
+}