Add support for GO batch delimiter in SQL Server

- per documentation, "not a statement" but acts like one in all other regards
- since it's a batch delimiter and statements can't extend beyond a batch, it also acts as a statement delimiter

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/keywords.rs

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

src/parser/mod.rs

@@ -512,6 +512,12 @@ impl<'a> Parser<'a> {
                     if expecting_statement_delimiter && word.keyword == Keyword::END {
                         break;
                     }
+                    // Treat batch delimiter as an end of statement
+                    if expecting_statement_delimiter && dialect_of!(self is MsSqlDialect) {
+                        if let Some(Statement::Go(GoStatement { count: _ })) = stmts.last() {
+                            expecting_statement_delimiter = false;
+                        }
+                    }
                 }
                 _ => {}
             }
@@ -715,6 +721,7 @@ impl<'a> Parser<'a> {
                     self.parse_vacuum()
                 }
                 Keyword::RESET => self.parse_reset().map(Into::into),
+                Keyword::GO => self.parse_go(),
                 _ => self.expected("an SQL statement", next_token),
             },
             Token::LParen => {
@@ -19637,6 +19644,61 @@ impl<'a> Parser<'a> {
         }))
     }
 
+    /// Parse [Statement::Go]
+    fn parse_go(&mut self) -> Result<Statement, ParserError> {
+        // previous token should be a newline (skipping non-newline whitespace)
+        // see also, `previous_token`
+        let mut look_back_count = 2;
+        loop {
+            let prev_index = self.index.saturating_sub(look_back_count);
+            if prev_index == 0 {
+                break;
+            }
+            let prev_token = self.token_at(prev_index);
+            match prev_token.token {
+                Token::Whitespace(ref w) => match w {
+                    Whitespace::Newline => break,
+                    _ => look_back_count += 1,
+                },
+                _ => {
+                    if prev_token == self.get_current_token() {
+                        // if we are at the start of the statement, we can skip this check
+                        break;
+                    }
+
+                    self.expected("newline before GO", prev_token.clone())?
+                }
+            };
+        }
+
+        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)?,
+            };
+        };
+
+        if self.peek_token().token == Token::SemiColon {
+            parser_err!(
+                "GO may not end with a semicolon",
+                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

tests/sqlparser_mssql.rs

@@ -2854,3 +2854,73 @@ fn parse_mssql_update_with_output_into() {
         "UPDATE employees SET salary = salary * 1.1 OUTPUT INSERTED.id, DELETED.salary, INSERTED.salary INTO @changes WHERE department = 'Engineering'",
     );
 }
+
+#[test]
+fn parse_mssql_go_keyword() {
+    let single_go_keyword = "USE some_database;\nGO";
+    let stmts = ms().parse_sql_statements(single_go_keyword).unwrap();
+    assert_eq!(stmts.len(), 2);
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }),);
+
+    let go_with_count = "SELECT 1;\nGO 5";
+    let stmts = ms().parse_sql_statements(go_with_count).unwrap();
+    assert_eq!(stmts.len(), 2);
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(5) }));
+
+    let bare_go = "GO";
+    let stmts = ms().parse_sql_statements(bare_go).unwrap();
+    assert_eq!(stmts.len(), 1);
+    assert_eq!(stmts[0], Statement::Go(GoStatement { count: None }));
+
+    let go_then_statements = "/* whitespace */ GO\nRAISERROR('This is a test', 16, 1);";
+    let stmts = ms().parse_sql_statements(go_then_statements).unwrap();
+    assert_eq!(stmts.len(), 2);
+    assert_eq!(stmts[0], Statement::Go(GoStatement { count: None }));
+    assert_eq!(
+        stmts[1],
+        Statement::RaisError {
+            message: Box::new(Expr::Value(
+                (Value::SingleQuotedString("This is a test".to_string())).with_empty_span()
+            )),
+            severity: Box::new(Expr::Value(number("16").with_empty_span())),
+            state: Box::new(Expr::Value(number("1").with_empty_span())),
+            arguments: vec![],
+            options: vec![],
+        }
+    );
+
+    let multiple_gos = "SELECT 1;\nGO 5\nSELECT 2;\n  GO";
+    let stmts = ms().parse_sql_statements(multiple_gos).unwrap();
+    assert_eq!(stmts.len(), 4);
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(5) }));
+    assert_eq!(stmts[3], Statement::Go(GoStatement { count: None }));
+
+    let comment_following_go = "USE some_database;\nGO -- okay";
+    let stmts = ms().parse_sql_statements(comment_following_go).unwrap();
+    assert_eq!(stmts.len(), 2);
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
+
+    let actually_column_alias = "SELECT NULL AS GO";
+    let stmt = ms().verified_only_select(actually_column_alias);
+    assert_eq!(
+        only(stmt.projection),
+        SelectItem::ExprWithAlias {
+            expr: Expr::Value(Value::Null.with_empty_span()),
+            alias: Ident::new("GO"),
+        }
+    );
+
+    let invalid_go_position = "SELECT 1; GO";
+    let err = ms().parse_sql_statements(invalid_go_position);
+    assert_eq!(
+        err.unwrap_err().to_string(),
+        "sql parser error: Expected: newline before GO, found: ;"
+    );
+
+    let invalid_go_count = "SELECT 1\nGO x";
+    let err = ms().parse_sql_statements(invalid_go_count);
+    assert_eq!(
+        err.unwrap_err().to_string(),
+        "sql parser error: Expected: end of statement, found: x"
+    );
+}