Implement ASSERT statement (apache#226)

Dandandan · web-flow · commit c24b0e01db2b · 2020-07-16T17:28:03.000+02:00
As supported by PostgreSQL and BigQuery (with some differences between them)
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -551,6 +551,14 @@ pub enum Statement {
     Rollback { chain: bool },
     /// CREATE SCHEMA
     CreateSchema { schema_name: ObjectName },
+
+    /// ASSERT <condition> [AS <message>]
+    Assert {
+        condition: Expr,
+        // AS or ,
+        separator: String,
+        message: Option<Expr>,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -810,6 +818,18 @@ impl fmt::Display for Statement {
                 write!(f, "ROLLBACK{}", if *chain { " AND CHAIN" } else { "" },)
             }
             Statement::CreateSchema { schema_name } => write!(f, "CREATE SCHEMA {}", schema_name),
+            Statement::Assert {
+                condition,
+                separator,
+                message,
+            } => {
+                write!(f, "ASSERT {}", condition)?;
+
+                if let Some(m) = message {
+                    write!(f, " {} {}", separator, m)?;
+                }
+                Ok(())
+            }
         }
     }
 }
diff --git a/src/dialect/keywords.rs b/src/dialect/keywords.rs
@@ -79,6 +79,7 @@ define_keywords!(
     AS,
     ASC,
     ASENSITIVE,
+    ASSERT,
     ASYMMETRIC,
     AT,
     ATOMIC,
diff --git a/src/parser.rs b/src/parser.rs
@@ -148,6 +148,7 @@ impl Parser {
                 Keyword::BEGIN => Ok(self.parse_begin()?),
                 Keyword::COMMIT => Ok(self.parse_commit()?),
                 Keyword::ROLLBACK => Ok(self.parse_rollback()?),
+                Keyword::ASSERT => Ok(self.parse_assert()?),
                 _ => self.expected("an SQL statement", Token::Word(w)),
             },
             Token::LParen => {
@@ -179,6 +180,22 @@ impl Parser {
         }
         Ok(expr)
     }
+    pub fn parse_assert(&mut self) -> Result<Statement, ParserError> {
+        let condition = self.parse_expr()?;
+        let (separator, message) = if self.consume_token(&Token::Comma) {
+            (",".to_string(), Some(self.parse_expr()?))
+        } else if self.parse_keyword(Keyword::AS) {
+            ("AS".to_string(), Some(self.parse_expr()?))
+        } else {
+            ("".to_string(), None)
+        };
+
+        Ok(Statement::Assert {
+            condition,
+            separator,
+            message,
+        })
+    }
 
     /// Parse an expression prefix
     pub fn parse_prefix(&mut self) -> Result<Expr, ParserError> {
diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -1156,6 +1156,46 @@ fn parse_create_table_with_multiple_on_delete_fails() {
     .expect_err("should have failed");
 }
 
+#[test]
+fn parse_assert() {
+    let sql = "ASSERT (SELECT COUNT(*) FROM table) > 0";
+    let ast = one_statement_parses_to(sql, "ASSERT (SELECT COUNT(*) FROM table) > 0");
+    match ast {
+        Statement::Assert {
+            condition: _condition,
+            separator,
+            message,
+        } => {
+            assert_eq!(message, None);
+            assert_eq!(separator, "");
+        }
+        _ => unreachable!(),
+    }
+}
+
+#[test]
+fn parse_assert_message() {
+    let sql = "ASSERT (SELECT COUNT(*) FROM table) > 0 AS 'No rows in table'";
+    let ast = one_statement_parses_to(
+        sql,
+        "ASSERT (SELECT COUNT(*) FROM table) > 0 AS 'No rows in table'",
+    );
+    match ast {
+        Statement::Assert {
+            condition: _condition,
+            message: Some(message),
+            separator,
+        } => {
+            assert_eq!(separator, "AS");
+            match message {
+                Expr::Value(Value::SingleQuotedString(s)) => assert_eq!(s, "No rows in table"),
+                _ => unreachable!(),
+            };
+        }
+        _ => unreachable!(),
+    }
+}
+
 #[test]
 fn parse_create_schema() {
     let sql = "CREATE SCHEMA X";
