MSSQL: Add support for WAITFOR statement

Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com>

Author: guan404ming

src/ast/mod.rs

@@ -4764,6 +4764,10 @@ pub enum Statement {
     ///
     /// See: <https://learn.microsoft.com/en-us/sql/t-sql/statements/print-transact-sql>
     Print(PrintStatement),
+    /// MSSQL `WAITFOR` statement.
+    ///
+    /// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql>
+    WaitFor(WaitForStatement),
     /// ```sql
     /// RETURN [ expression ]
     /// ```
@@ -6217,6 +6221,7 @@ impl fmt::Display for Statement {
             }
             Statement::Throw(s) => write!(f, "{s}"),
             Statement::Print(s) => write!(f, "{s}"),
+            Statement::WaitFor(s) => write!(f, "{s}"),
             Statement::Return(r) => write!(f, "{r}"),
             Statement::List(command) => write!(f, "LIST {command}"),
             Statement::Remove(command) => write!(f, "REMOVE {command}"),
@@ -10964,6 +10969,47 @@ impl fmt::Display for PrintStatement {
     }
 }
 
+/// The type of `WAITFOR` statement (MSSQL).
+///
+/// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum WaitForType {
+    /// `WAITFOR DELAY 'time_to_pass'`
+    Delay,
+    /// `WAITFOR TIME 'time_to_execute'`
+    Time,
+}
+
+impl fmt::Display for WaitForType {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            WaitForType::Delay => write!(f, "DELAY"),
+            WaitForType::Time => write!(f, "TIME"),
+        }
+    }
+}
+
+/// MSSQL `WAITFOR` statement.
+///
+/// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct WaitForStatement {
+    /// `DELAY` or `TIME`.
+    pub wait_type: WaitForType,
+    /// The time expression.
+    pub expr: Expr,
+}
+
+impl fmt::Display for WaitForStatement {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "WAITFOR {} {}", self.wait_type, self.expr)
+    }
+}
+
 /// Represents a `Return` statement.
 ///
 /// [MsSql triggers](https://learn.microsoft.com/en-us/sql/t-sql/statements/create-trigger-transact-sql)

src/ast/spans.rs

@@ -483,6 +483,7 @@ impl Spanned for Statement {
             Statement::RaisError { .. } => Span::empty(),
             Statement::Throw(_) => Span::empty(),
             Statement::Print { .. } => Span::empty(),
+            Statement::WaitFor(_) => Span::empty(),
             Statement::Return { .. } => Span::empty(),
             Statement::List(..) | Statement::Remove(..) => Span::empty(),
             Statement::ExportData(ExportData {

src/dialect/mssql.rs

@@ -18,7 +18,7 @@
 use crate::ast::helpers::attached_token::AttachedToken;
 use crate::ast::{
     BeginEndStatements, ConditionalStatementBlock, ConditionalStatements, CreateTrigger,
-    GranteesType, IfStatement, Statement,
+    GranteesType, IfStatement, Statement, WaitForStatement, WaitForType,
 };
 use crate::dialect::Dialect;
 use crate::keywords::Keyword;
@@ -183,6 +183,8 @@ impl Dialect for MsSqlDialect {
             Keyword::TRIGGER,
         ]) {
             Some(self.parse_create_trigger(parser, true))
+        } else if parser.parse_keyword(Keyword::WAITFOR) {
+            Some(self.parse_waitfor(parser))
         } else {
             None
         }
@@ -347,4 +349,19 @@ impl MsSqlDialect {
         }
         Ok(stmts)
     }
+
+    /// Parse `WAITFOR { DELAY 'time' | TIME 'time' }`
+    ///
+    /// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql>
+    fn parse_waitfor(&self, parser: &mut Parser) -> Result<Statement, ParserError> {
+        let wait_type = if parser.parse_keyword(Keyword::DELAY) {
+            WaitForType::Delay
+        } else if parser.parse_keyword(Keyword::TIME) {
+            WaitForType::Time
+        } else {
+            return parser.expected("DELAY or TIME", parser.peek_token());
+        };
+        let expr = parser.parse_expr()?;
+        Ok(Statement::WaitFor(WaitForStatement { wait_type, expr }))
+    }
 }

src/keywords.rs

@@ -312,6 +312,7 @@ define_keywords!(
     DEFINE,
     DEFINED,
     DEFINER,
+    DELAY,
     DELAYED,
     DELAY_KEY_WRITE,
     DELEGATED,
@@ -1138,6 +1139,7 @@ define_keywords!(
     VIRTUAL,
     VOLATILE,
     VOLUME,
+    WAITFOR,
     WAREHOUSE,
     WAREHOUSES,
     WEEK,

tests/sqlparser_mssql.rs

@@ -1702,6 +1702,43 @@ fn test_parse_throw() {
     );
 }
 
+#[test]
+fn test_parse_waitfor() {
+    // WAITFOR DELAY
+    let sql = "WAITFOR DELAY '00:00:05'";
+    let stmt = ms().verified_stmt(sql);
+    assert_eq!(
+        stmt,
+        Statement::WaitFor(WaitForStatement {
+            wait_type: WaitForType::Delay,
+            expr: Expr::Value(
+                (Value::SingleQuotedString("00:00:05".to_string())).with_empty_span()
+            ),
+        })
+    );
+
+    // WAITFOR TIME
+    let sql = "WAITFOR TIME '14:30:00'";
+    let stmt = ms().verified_stmt(sql);
+    assert_eq!(
+        stmt,
+        Statement::WaitFor(WaitForStatement {
+            wait_type: WaitForType::Time,
+            expr: Expr::Value(
+                (Value::SingleQuotedString("14:30:00".to_string())).with_empty_span()
+            ),
+        })
+    );
+
+    // WAITFOR DELAY with variable
+    let sql = "WAITFOR DELAY @WaitTime";
+    let _ = ms().verified_stmt(sql);
+
+    // Error: WAITFOR without DELAY or TIME
+    let res = ms().parse_sql_statements("WAITFOR '00:00:05'");
+    assert!(res.is_err());
+}
+
 #[test]
 fn parse_use() {
     let valid_object_names = [