Fix not treating GO as a statement delimiter next to a query

Author: aharpervc

src/dialect/mssql.rs

@@ -23,7 +23,7 @@ use crate::ast::{
 use crate::dialect::Dialect;
 use crate::keywords::{self, Keyword};
 use crate::parser::{Parser, ParserError};
-use crate::tokenizer::Token;
+use crate::tokenizer::{Token, Whitespace};
 #[cfg(not(feature = "std"))]
 use alloc::{vec, vec::Vec};
 
@@ -128,7 +128,29 @@ impl Dialect for MsSqlDialect {
         &[GranteesType::Public]
     }
 
-    fn is_column_alias(&self, kw: &Keyword, _parser: &mut Parser) -> bool {
+    fn is_column_alias(&self, kw: &Keyword, parser: &mut Parser) -> bool {
+        // if:
+        // - keyword is `GO`, and
+        // - looking backwards there's only (any) whitespace preceded by a newline
+        // then: `GO` iSN'T a column alias
+        if kw == &Keyword::GO {
+            let mut look_back_count = 2;
+            loop {
+                let prev_index = parser.index().saturating_sub(look_back_count);
+                if prev_index == 0 {
+                    break;
+                }
+                let prev_token = parser.token_at(prev_index);
+                match prev_token.token {
+                    Token::Whitespace(ref w) => match w {
+                        Whitespace::Newline => return false,
+                        _ => look_back_count += 1,
+                    },
+                    _ => break,
+                };
+            }
+        }
+
         !keywords::RESERVED_FOR_COLUMN_ALIAS.contains(kw) && !RESERVED_FOR_COLUMN_ALIAS.contains(kw)
     }
 

src/parser/mod.rs

@@ -483,6 +483,10 @@ impl<'a> Parser<'a> {
                     if expecting_statement_delimiter && word.keyword == Keyword::END {
                         break;
                     }
+
+                    if expecting_statement_delimiter && word.keyword == Keyword::GO {
+                        expecting_statement_delimiter = false;
+                    }
                 }
                 _ => {}
             }

tests/sqlparser_mssql.rs

@@ -2520,6 +2520,11 @@ fn parse_mssql_go_keyword() {
     assert_eq!(stmts.len(), 2);
     assert_eq!(stmts[1], Statement::Go(GoStatement { count: Some(5) }));
 
+    let go_statement_delimiter = "SELECT 1\nGO";
+    let stmts = ms().parse_sql_statements(go_statement_delimiter).unwrap();
+    assert_eq!(stmts.len(), 2);
+    assert_eq!(stmts[1], Statement::Go(GoStatement { count: None }));
+
     let bare_go = "GO";
     let stmts = ms().parse_sql_statements(bare_go).unwrap();
     assert_eq!(stmts.len(), 1);
@@ -2553,15 +2558,22 @@ fn parse_mssql_go_keyword() {
     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 actually_column_alias = "SELECT NULL GO";
+    let stmts = ms().parse_sql_statements(actually_column_alias).unwrap();
+    assert_eq!(stmts.len(), 1);
+    match &stmts[0] {
+        Statement::Query(query) => {
+            let select = query.body.as_select().unwrap();
+            assert_eq!(
+                only(select.clone().projection),
+                SelectItem::ExprWithAlias {
+                    expr: Expr::Value(Value::Null.with_empty_span()),
+                    alias: Ident::new("GO"),
+                }
+            );
         }
-    );
+        _ => panic!("Expected Query statement"),
+    }
 
     let invalid_go_position = "SELECT 1; GO";
     let err = ms().parse_sql_statements(invalid_go_position);
@@ -2574,7 +2586,7 @@ fn parse_mssql_go_keyword() {
     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"
+        "sql parser error: Expected: literal int or newline, found: x"
     );
 }