From 525ba527bb3870e07617e7388bed8750820cb185 Mon Sep 17 00:00:00 2001
From: Maciej Obuchowski <obuchowski.maciej@gmail.com>
Date: Fri, 29 Apr 2022 20:11:11 +0200
Subject: [PATCH] snowflake: add qualify expression (#465)

Signed-off-by: Maciej Obuchowski <obuchowski.maciej@gmail.com>
---
 src/ast/query.rs              |  5 +++++
 src/keywords.rs               |  2 ++
 src/parser.rs                 |  7 ++++++
 tests/sqlparser_common.rs     | 42 ++++++++++++++++++++++++++++++++++-
 tests/sqlparser_mysql.rs      |  2 ++
 tests/sqlparser_postgres.rs   |  1 +
 tests/sqpparser_clickhouse.rs |  3 ++-
 7 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/src/ast/query.rs b/src/ast/query.rs
index 82341f63..7526c23c 100644
--- a/src/ast/query.rs
+++ b/src/ast/query.rs
@@ -154,6 +154,8 @@ pub struct Select {
     pub sort_by: Vec<Expr>,
     /// HAVING
     pub having: Option<Expr>,
+    /// QUALIFY (Snowflake)
+    pub qualify: Option<Expr>,
 }
 
 impl fmt::Display for Select {
@@ -202,6 +204,9 @@ impl fmt::Display for Select {
         if let Some(ref having) = self.having {
             write!(f, " HAVING {}", having)?;
         }
+        if let Some(ref qualify) = self.qualify {
+            write!(f, " QUALIFY {}", qualify)?;
+        }
         Ok(())
     }
 }
diff --git a/src/keywords.rs b/src/keywords.rs
index ad31bda5..ccb4f3f9 100644
--- a/src/keywords.rs
+++ b/src/keywords.rs
@@ -382,6 +382,7 @@ define_keywords!(
     PROCEDURE,
     PROGRAM,
     PURGE,
+    QUALIFY,
     QUARTER,
     QUOTE,
     RANGE,
@@ -584,6 +585,7 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     // for MSSQL-specific OUTER APPLY (seems reserved in most dialects)
     Keyword::OUTER,
     Keyword::SET,
+    Keyword::QUALIFY,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`
diff --git a/src/parser.rs b/src/parser.rs
index 1ca493b9..e855bdf4 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -3167,6 +3167,12 @@ impl<'a> Parser<'a> {
             None
         };
 
+        let qualify = if self.parse_keyword(Keyword::QUALIFY) {
+            Some(self.parse_expr()?)
+        } else {
+            None
+        };
+
         Ok(Select {
             distinct,
             top,
@@ -3180,6 +3186,7 @@ impl<'a> Parser<'a> {
             distribute_by,
             sort_by,
             having,
+            qualify,
         })
     }
 
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
index 8471b2c5..e0c7f34b 100644
--- a/tests/sqlparser_common.rs
+++ b/tests/sqlparser_common.rs
@@ -1322,6 +1322,45 @@ fn parse_select_having() {
     assert!(select.having.is_some());
 }
 
+#[cfg(feature = "bigdecimal")]
+#[test]
+fn parse_select_qualify() {
+    let sql = "SELECT i, p, o FROM qt QUALIFY ROW_NUMBER() OVER (PARTITION BY p ORDER BY o) = 1";
+    let select = verified_only_select(sql);
+    assert_eq!(
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Function(Function {
+                name: ObjectName(vec![Ident::new("ROW_NUMBER")]),
+                args: vec![],
+                over: Some(WindowSpec {
+                    partition_by: vec![Expr::Identifier(Ident::new("p"))],
+                    order_by: vec![OrderByExpr {
+                        expr: Expr::Identifier(Ident::new("o")),
+                        asc: None,
+                        nulls_first: None
+                    }],
+                    window_frame: None
+                }),
+                distinct: false
+            })),
+            op: BinaryOperator::Eq,
+            right: Box::new(Expr::Value(number("1")))
+        }),
+        select.qualify
+    );
+
+    let sql = "SELECT i, p, o, ROW_NUMBER() OVER (PARTITION BY p ORDER BY o) AS row_num FROM qt QUALIFY row_num = 1";
+    let select = verified_only_select(sql);
+    assert_eq!(
+        Some(Expr::BinaryOp {
+            left: Box::new(Expr::Identifier(Ident::new("row_num"))),
+            op: BinaryOperator::Eq,
+            right: Box::new(Expr::Value(number("1")))
+        }),
+        select.qualify
+    );
+}
+
 #[test]
 fn parse_limit_accepts_all() {
     one_statement_parses_to(
@@ -4336,7 +4375,8 @@ fn parse_merge() {
                         cluster_by: vec![],
                         distribute_by: vec![],
                         sort_by: vec![],
-                        having: None
+                        having: None,
+                        qualify: None
                     })),
                     order_by: vec![],
                     limit: None,
diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
index e3566bca..454c6639 100644
--- a/tests/sqlparser_mysql.rs
+++ b/tests/sqlparser_mysql.rs
@@ -314,6 +314,7 @@ fn parse_quote_identifiers_2() {
                 distribute_by: vec![],
                 sort_by: vec![],
                 having: None,
+                qualify: None
             })),
             order_by: vec![],
             limit: None,
@@ -754,6 +755,7 @@ fn parse_substring_in_select() {
                         distribute_by: vec![],
                         sort_by: vec![],
                         having: None,
+                        qualify: None
                     })),
                     order_by: vec![],
                     limit: None,
diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
index eb42edc8..125ce9f6 100644
--- a/tests/sqlparser_postgres.rs
+++ b/tests/sqlparser_postgres.rs
@@ -450,6 +450,7 @@ fn parse_update_set_from() {
                             distribute_by: vec![],
                             sort_by: vec![],
                             having: None,
+                            qualify: None
                         })),
                         order_by: vec![],
                         limit: None,
diff --git a/tests/sqpparser_clickhouse.rs b/tests/sqpparser_clickhouse.rs
index dfd55520..59fc91c8 100644
--- a/tests/sqpparser_clickhouse.rs
+++ b/tests/sqpparser_clickhouse.rs
@@ -96,7 +96,8 @@ fn parse_map_access_expr() {
             cluster_by: vec![],
             distribute_by: vec![],
             sort_by: vec![],
-            having: None
+            having: None,
+            qualify: None
         },
         select
     );