Extract common logic to look back for a newline to a helper

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, Whitespace};
+use crate::tokenizer::Token;
 #[cfg(not(feature = "std"))]
 use alloc::{vec, vec::Vec};
 
@@ -129,26 +129,14 @@ impl Dialect for MsSqlDialect {
     }
 
     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,
-                };
-            }
+        // 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
+                .expect_previously_only_whitespace_until_newline()
+                .is_ok()
+        {
+            return false;
         }
 
         !keywords::RESERVED_FOR_COLUMN_ALIAS.contains(kw) && !RESERVED_FOR_COLUMN_ALIAS.contains(kw)

src/parser/mod.rs

@@ -4150,6 +4150,44 @@ impl<'a> Parser<'a> {
         )
     }
 
+    /// Look backwards in the token stream and expect that there was only whitespace tokens until the previous newline
+    pub fn expect_previously_only_whitespace_until_newline(&mut self) -> Result<(), ParserError> {
+        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,
+                    // special consideration required for single line comments since that string includes the newline
+                    Whitespace::SingleLineComment { comment, prefix: _ } => {
+                        if comment.ends_with('\n') {
+                            break;
+                        }
+                        look_back_count += 1;
+                    }
+                    _ => look_back_count += 1,
+                },
+                _ => {
+                    let current_token = self.get_current_token();
+                    if prev_token == current_token {
+                        // if we are at the start of the statement, we can skip this check
+                        break;
+                    }
+
+                    self.expected(
+                        &format!("newline before current token ({})", current_token),
+                        prev_token.clone(),
+                    )?
+                }
+            };
+        }
+        Ok(())
+    }
+
     /// If the current token is the `expected` keyword, consume it and returns
     /// true. Otherwise, no tokens are consumed and returns false.
     #[must_use]
@@ -16395,36 +16433,7 @@ 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,
-                    Whitespace::SingleLineComment { comment, prefix: _ } => {
-                        if comment.ends_with('\n') {
-                            break;
-                        }
-                        look_back_count += 1;
-                    }
-                    _ => 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())?
-                }
-            };
-        }
+        self.expect_previously_only_whitespace_until_newline()?;
 
         let count = loop {
             // using this peek function because we want to halt this statement parsing upon newline

tests/sqlparser_mssql.rs

@@ -2603,7 +2603,7 @@ fn parse_mssql_go_keyword() {
     let err = ms().parse_sql_statements(invalid_go_position);
     assert_eq!(
         err.unwrap_err().to_string(),
-        "sql parser error: Expected: newline before GO, found: ;"
+        "sql parser error: Expected: newline before current token (GO), found: ;"
     );
 
     let invalid_go_count = "SELECT 1\nGO x";