squash-merge mssql-go-keyword

Author: aharpervc

src/ast/mod.rs

@@ -4833,6 +4833,12 @@ pub enum Statement {
     /// ```
     /// [PostgreSQL](https://www.postgresql.org/docs/current/sql-reset.html)
     Reset(ResetStatement),
+    /// Go (MsSql)
+    ///
+    /// GO is not a Transact-SQL statement; it is a command recognized by various tools as a batch delimiter
+    ///
+    /// See: <https://learn.microsoft.com/en-us/sql/t-sql/language-elements/sql-server-utilities-statements-go>
+    Go(GoStatement),
 }
 
 impl From<Analyze> for Statement {
@@ -6247,6 +6253,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::Go(s) => write!(f, "{s}"),
             Statement::Return(r) => write!(f, "{r}"),
             Statement::List(command) => write!(f, "LIST {command}"),
             Statement::Remove(command) => write!(f, "REMOVE {command}"),
@@ -11547,6 +11554,27 @@ impl fmt::Display for CreateTableLikeDefaults {
     }
 }
 
+/// Represents a `GO` statement.
+///
+/// [MsSql](https://learn.microsoft.com/en-us/sql/t-sql/language-elements/sql-server-utilities-statements-go)
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct GoStatement {
+    /// How many times the batch should be executed, if specified (e.g., `GO 10`).
+    pub count: Option<u64>,
+}
+
+impl Display for GoStatement {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        if let Some(count) = self.count {
+            write!(f, "GO {count}")
+        } else {
+            write!(f, "GO")
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/ast/spans.rs

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

src/dialect/mssql.rs

@@ -132,6 +132,12 @@ impl Dialect for MsSqlDialect {
     }
 
     fn is_select_item_alias(&self, explicit: bool, kw: &Keyword, parser: &mut Parser) -> bool {
+        // if we find maybe whitespace then a newline looking backward, then `GO` ISN'T a column alias
+        // if we can't find a newline then we assume that `GO` IS a column alias
+        if kw == &Keyword::GO && parser.prev_only_whitespace_until_newline() {
+            return false;
+        }
+
         match kw {
             // List of keywords that cannot be used as select item (column) aliases in MSSQL
             // regardless of whether the alias is explicit or implicit.
@@ -188,6 +194,7 @@ impl Dialect for MsSqlDialect {
             | Keyword::PRINT
             | Keyword::WHILE
             | Keyword::RETURN
+            | Keyword::GO
             | Keyword::THROW
             | Keyword::RAISERROR
             | Keyword::MERGE => false,

src/keywords.rs

@@ -463,6 +463,7 @@ define_keywords!(
     GIN,
     GIST,
     GLOBAL,
+    GO,
     GRANT,
     GRANTED,
     GRANTS,

src/parser/mod.rs

@@ -531,6 +531,12 @@ impl<'a> Parser<'a> {
                     if expecting_statement_delimiter && word.keyword == Keyword::END {
                         break;
                     }
+
+                    // MSSQL: the `GO` keyword is a batch separator which also means it concludes the current statement
+                    // `GO` may not be followed by a semicolon, so turn off that expectation
+                    if expecting_statement_delimiter && word.keyword == Keyword::GO {
+                        expecting_statement_delimiter = false;
+                    }
                 }
                 // don't expect a semicolon statement delimiter after a newline when not otherwise required
                 Token::Whitespace(Whitespace::Newline) => {
@@ -546,8 +552,10 @@ impl<'a> Parser<'a> {
             }
 
             let statement = self.parse_statement()?;
+            // MSSQL: the `GO` keyword is a batch separator which also means it concludes the current statement
+            // `GO` may not be followed by a semicolon, so turn off that expectation
+            expecting_statement_delimiter = !matches!(statement, Statement::Go(_)) && self.options.require_semicolon_stmt_delimiter;
             stmts.push(statement);
-            expecting_statement_delimiter = self.options.require_semicolon_stmt_delimiter;
         }
         Ok(stmts)
     }
@@ -740,6 +748,10 @@ impl<'a> Parser<'a> {
                     self.parse_vacuum()
                 }
                 Keyword::RESET => self.parse_reset().map(Into::into),
+                Keyword::GO => {
+                    self.prev_token();
+                    self.parse_go()
+                }
                 _ => self.expected("an SQL statement", next_token),
             },
             Token::LParen => {
@@ -4473,6 +4485,17 @@ impl<'a> Parser<'a> {
         self.tokens.get(self.index + n).unwrap_or(&EOF_TOKEN)
     }
 
+    /// Return nth previous token, possibly whitespace
+    /// (or [`Token::EOF`] when before the beginning of the stream).
+    pub(crate) fn peek_prev_nth_token_no_skip_ref(&self, n: usize) -> &TokenWithSpan {
+        // 0 = next token, -1 = current token, -2 = previous token
+        let peek_index = self.index.saturating_sub(1).saturating_sub(n);
+        if peek_index == 0 {
+            return &EOF_TOKEN;
+        }
+        self.tokens.get(peek_index).unwrap_or(&EOF_TOKEN)
+    }
+
     /// Return true if the next tokens exactly `expected`
     ///
     /// Does not advance the current token.
@@ -4589,6 +4612,29 @@ impl<'a> Parser<'a> {
         )
     }
 
+    /// Look backwards in the token stream and expect that there was only whitespace tokens until the previous newline or beginning of string
+    pub(crate) fn prev_only_whitespace_until_newline(&mut self) -> bool {
+        let mut look_back_count = 1;
+        loop {
+            let prev_token = self.peek_prev_nth_token_no_skip_ref(look_back_count);
+            match prev_token.token {
+                Token::EOF => break true,
+                Token::Whitespace(ref w) => match w {
+                    Whitespace::Newline => break true,
+                    // special consideration required for single line comments since that string includes the newline
+                    Whitespace::SingleLineComment { comment, prefix: _ } => {
+                        if comment.ends_with('\n') {
+                            break true;
+                        }
+                        look_back_count += 1;
+                    }
+                    _ => look_back_count += 1,
+                },
+                _ => break false,
+            };
+        }
+    }
+
     /// If the current token is the `expected` keyword, consume it and returns
     /// true. Otherwise, no tokens are consumed and returns false.
     #[must_use]
@@ -19700,6 +19746,71 @@ impl<'a> Parser<'a> {
         }))
     }
 
+    /// Parse [Statement::Go]
+    fn parse_go(&mut self) -> Result<Statement, ParserError> {
+        self.expect_keyword_is(Keyword::GO)?;
+
+        // disambiguate between GO as batch delimiter & GO as identifier (etc)
+        // compare:
+        // ```sql
+        // select 1 go
+        // ```
+        // vs
+        // ```sql
+        // select 1
+        // go
+        // ```
+        if !self.prev_only_whitespace_until_newline() {
+            parser_err!(
+                "GO may only be preceded by whitespace on a line",
+                self.peek_token().span.start
+            )?;
+        }
+
+        let count = loop {
+            // using this peek function because we want to halt this statement parsing upon newline
+            let next_token = self.peek_token_no_skip();
+            match next_token.token {
+                Token::EOF => break None::<u64>,
+                Token::Whitespace(ref w) => match w {
+                    Whitespace::Newline => break None,
+                    _ => _ = self.next_token_no_skip(),
+                },
+                Token::Number(s, _) => {
+                    let value = Some(Self::parse::<u64>(s, next_token.span.start)?);
+                    self.advance_token();
+                    break value;
+                }
+                _ => self.expected("literal int or newline", next_token)?,
+            };
+        };
+
+        loop {
+            let next_token = self.peek_token_no_skip();
+            match next_token.token {
+                Token::EOF => break,
+                Token::Whitespace(ref w) => match w {
+                    Whitespace::Newline => break,
+                    Whitespace::SingleLineComment { comment, prefix: _ } => {
+                        if comment.ends_with('\n') {
+                            break;
+                        }
+                        _ = self.next_token_no_skip();
+                    }
+                    _ => _ = self.next_token_no_skip(),
+                },
+                _ => {
+                    parser_err!(
+                        "GO must be followed by a newline or EOF",
+                        self.peek_token().span.start
+                    )?;
+                }
+            };
+        }
+
+        Ok(Statement::Go(GoStatement { count }))
+    }
+
     /// Consume the parser and return its underlying token buffer
     pub fn into_tokens(self) -> Vec<TokenWithSpan> {
         self.tokens
@@ -20066,6 +20177,31 @@ mod tests {
         })
     }
 
+    #[test]
+    fn test_peek_prev_nth_token_no_skip_ref() {
+        all_dialects().run_parser_method(
+            "SELECT 1;\n-- a comment\nRAISERROR('test', 16, 0);",
+            |parser| {
+                parser.index = 1;
+                assert_eq!(parser.peek_prev_nth_token_no_skip_ref(0), &Token::EOF);
+                assert_eq!(parser.index, 1);
+                parser.index = 7;
+                assert_eq!(
+                    parser.token_at(parser.index - 1).token,
+                    Token::Word(Word {
+                        value: "RAISERROR".to_string(),
+                        quote_style: None,
+                        keyword: Keyword::RAISERROR,
+                    })
+                );
+                assert_eq!(
+                    parser.peek_prev_nth_token_no_skip_ref(2),
+                    &Token::Whitespace(Whitespace::Newline)
+                );
+            },
+        );
+    }
+
     #[cfg(test)]
     mod test_parse_data_type {
         use crate::ast::{

src/test_utils.rs

@@ -153,7 +153,7 @@ impl TestedDialects {
     /// 2. re-serializing the result of parsing `sql` produces the same
     ///    `canonical` sql string
     ///
-    ///  For multiple statements, use [`statements_parse_to`].
+    /// For multiple statements, use [`statements_parse_to`].
     pub fn one_statement_parses_to(&self, sql: &str, canonical: &str) -> Statement {
         let mut statements = self.parse_sql_statements(sql).expect(sql);
         assert_eq!(statements.len(), 1);
@@ -170,8 +170,15 @@ impl TestedDialects {
     }
 
     /// The same as [`one_statement_parses_to`] but it works for a multiple statements
-    pub fn statements_parse_to(&self, sql: &str, canonical: &str) -> Vec<Statement> {
+    pub fn statements_parse_to(
+        &self,
+        sql: &str,
+        statement_count: usize,
+        canonical: &str,
+    ) -> Vec<Statement> {
         let statements = self.parse_sql_statements(sql).expect(sql);
+        assert_eq!(statements.len(), statement_count);
+
         if !canonical.is_empty() && sql != canonical {
             assert_eq!(self.parse_sql_statements(canonical).unwrap(), statements);
         } else {