Move WAITFOR parsing to generic parser

Co-authored-by: Ifeanyi Ubah <7816405+iffyio@users.noreply.github.com>
Signed-off-by: Guan-Ming Chiu <guanmingchiu@gmail.com>

Author: guan404ming

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, WaitForStatement, WaitForType,
+    GranteesType, IfStatement, Statement,
 };
 use crate::dialect::Dialect;
 use crate::keywords::Keyword;
@@ -183,8 +183,6 @@ 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
         }
@@ -349,19 +347,4 @@ 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/parser/mod.rs

@@ -703,6 +703,8 @@ impl<'a> Parser<'a> {
                 // `COMMENT` is snowflake specific https://docs.snowflake.com/en/sql-reference/sql/comment
                 Keyword::COMMENT if self.dialect.supports_comment_on() => self.parse_comment(),
                 Keyword::PRINT => self.parse_print(),
+                // `WAITFOR` is MSSQL specific https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql
+                Keyword::WAITFOR => self.parse_waitfor(),
                 Keyword::RETURN => self.parse_return(),
                 Keyword::EXPORT => {
                     self.prev_token();
@@ -19289,6 +19291,21 @@ impl<'a> Parser<'a> {
         }))
     }
 
+    /// Parse [Statement::WaitFor]
+    ///
+    /// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/waitfor-transact-sql>
+    fn parse_waitfor(&mut self) -> Result<Statement, ParserError> {
+        let wait_type = if self.parse_keyword(Keyword::DELAY) {
+            WaitForType::Delay
+        } else if self.parse_keyword(Keyword::TIME) {
+            WaitForType::Time
+        } else {
+            return self.expected("DELAY or TIME", self.peek_token());
+        };
+        let expr = self.parse_expr()?;
+        Ok(Statement::WaitFor(WaitForStatement { wait_type, expr }))
+    }
+
     /// Parse [Statement::Return]
     fn parse_return(&mut self) -> Result<Statement, ParserError> {
         match self.maybe_parse(|p| p.parse_expr())? {

tests/sqlparser_mssql.rs

@@ -1706,7 +1706,7 @@ fn test_parse_throw() {
 fn test_parse_waitfor() {
     // WAITFOR DELAY
     let sql = "WAITFOR DELAY '00:00:05'";
-    let stmt = ms().verified_stmt(sql);
+    let stmt = ms_and_generic().verified_stmt(sql);
     assert_eq!(
         stmt,
         Statement::WaitFor(WaitForStatement {
@@ -1719,7 +1719,7 @@ fn test_parse_waitfor() {
 
     // WAITFOR TIME
     let sql = "WAITFOR TIME '14:30:00'";
-    let stmt = ms().verified_stmt(sql);
+    let stmt = ms_and_generic().verified_stmt(sql);
     assert_eq!(
         stmt,
         Statement::WaitFor(WaitForStatement {