fix(parser)!: treat a GO line inside a SQL Server statement as an error

CHANGELOG.md

@@ -159,6 +159,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   exist; the limits are constants.
 
 ### Behavior changes
+- SQL Server: a line holding only `GO [n]` inside a statement is now a parse
+  error ("GO batch separator inside a statement"), matching sqlcmd, instead of
+  a column or alias named `go`. `GO 0` / `GO 1.5` alone on a line report
+  "invalid GO repeat count".
 - A raw NUL or Ctrl-Z byte inside a string literal is an error in dialects
   whose literals cannot spell it back (NUL without backslash escapes; both
   without a dialect; Ctrl-Z in ClickHouse; both in Snowflake). Formatting used

pkg/gosqlx/batch_separator_test.go

@@ -81,3 +81,42 @@ func TestBatchSeparatorInvalidCounts(t *testing.T) {
 		})
 	}
 }
+
+// TestBatchSeparatorInvalidCountMessage checks that GO with an invalid repeat
+// count alone on its line is reported as such, not as a FROM / end of
+// statement error from the statement before it.
+func TestBatchSeparatorInvalidCountMessage(t *testing.T) {
+	for _, count := range []string{"0", "1.5", "999999999999999999999999999999"} {
+		t.Run(count, func(t *testing.T) {
+			tree, err := ParseWithDialect("SELECT 1\nGO "+count+"\nSELECT 2", keywords.DialectSQLServer)
+			if err == nil {
+				ast.ReleaseAST(tree)
+				t.Fatal("accepted invalid GO repeat count")
+			}
+			if !strings.Contains(err.Error(), "invalid GO repeat count") {
+				t.Fatalf("error = %v, want an invalid GO repeat count error", err)
+			}
+		})
+	}
+}
+
+// TestBatchSeparatorInsideStatementIsAnError checks that formatting refuses
+// a statement that runs across a GO line (sqlcmd would split it there)
+// instead of writing go as a column name, with or without comments.
+func TestBatchSeparatorInsideStatementIsAnError(t *testing.T) {
+	for _, sql := range []string{
+		"SELECT a,\ngo\nFROM t",
+		"SELECT a, -- c\ngo\nFROM t",
+		"SELECT a /* c */,\nGO 2\nFROM t",
+	} {
+		t.Run(sql, func(t *testing.T) {
+			out, err := FormatWithDialect(sql, DefaultFormatOptions(), keywords.DialectSQLServer)
+			if err == nil {
+				t.Fatalf("formatted as %q, want an error", out)
+			}
+			if !strings.Contains(err.Error(), "GO batch separator inside a statement") {
+				t.Fatalf("error = %v", err)
+			}
+		})
+	}
+}

pkg/sql/parser/parser.go

@@ -456,11 +456,15 @@ func (p *Parser) parseTokens(tokens []models.TokenWithSpan) (*ast.AST, error) {
 			continue
 		}
 
+		startPos := p.currentPos
 		stmt, err := p.parseStatement()
 		if err == nil {
 			result.Statements = append(result.Statements, stmt)
 			result.StatementStarts = append(result.StatementStarts, start)
 		}
+		if p.IsSQLServer() { // batchLineError returns err unchanged elsewhere
+			_, err = p.batchLineError(startPos, err)
+		}
 		if err != nil {
 			// Clean up the AST on error
 			ast.ReleaseAST(result)
@@ -664,11 +668,15 @@ func (p *Parser) parseContextTokens(ctx context.Context, tokens []models.TokenWi
 			continue
 		}
 
+		startPos := p.currentPos
 		stmt, err := p.parseStatement()
 		if err == nil {
 			result.Statements = append(result.Statements, stmt)
 			result.StatementStarts = append(result.StatementStarts, start)
 		}
+		if p.IsSQLServer() { // batchLineError returns err unchanged elsewhere
+			_, err = p.batchLineError(startPos, err)
+		}
 		if err != nil {
 			// Clean up the AST on error
 			ast.ReleaseAST(result)

pkg/sql/parser/recovery.go

@@ -207,6 +207,15 @@ func (p *Parser) parseWithRecovery(tokens []models.TokenWithSpan) ([]ast.Stateme
 
 		stmtStartPos := p.currentPos
 		stmt, err := p.parseStatement()
+		if err == nil {
+			// A statement that ran across a GO line is reported, and parsing
+			// resumes at that line so the batch boundary is kept.
+			if at, gerr := p.batchLineError(stmtStartPos, nil); gerr != nil {
+				errors = append(errors, gerr)
+				p.seek(at)
+				continue
+			}
+		}
 		if err != nil {
 			// Create a ParseError with position info, preserving original error
 			loc := p.currentLocation()

pkg/sql/parser/statement_boundary.go

@@ -24,6 +24,7 @@ package parser
 
 import (
 	"github.com/ajitpratap0/GoSQLX/internal/batch"
+	goerrors "github.com/ajitpratap0/GoSQLX/pkg/errors"
 	"github.com/ajitpratap0/GoSQLX/pkg/models"
 	"github.com/ajitpratap0/GoSQLX/pkg/sql/ast"
 )
@@ -62,6 +63,39 @@ func (p *Parser) parseBatchSeparatorAt() (*ast.BatchSeparator, bool) {
 	return sep, true
 }
 
+// batchLineError checks, in SQL Server, the tokens a statement starting at
+// token index start ran over (up to the current token when parseErr is set,
+// since the statement parser stopped there) for a line holding only GO [n].
+// sqlcmd ends the batch at such a line wherever it appears, so a statement
+// that runs across it is an error; GO with an invalid repeat count (GO 0,
+// GO 1.5) gets an error of its own. It returns parseErr when there is no
+// such line, otherwise the index of the GO token and the GO error.
+func (p *Parser) batchLineError(start int, parseErr error) (int, error) {
+	if !p.IsSQLServer() {
+		return -1, parseErr
+	}
+	last := p.currentPos - 1
+	if parseErr != nil {
+		last = p.currentPos
+	}
+	if last >= len(p.tokens) {
+		last = len(p.tokens) - 1
+	}
+	for i := start; i <= last; i++ {
+		if t := &p.tokens[i].Token; t.Type != models.TokenTypeIdentifier || len(t.Value) != 2 {
+			continue
+		}
+		if _, valid, ok := batch.Line(p.tokens, i); ok {
+			loc := p.tokens[i].Start
+			if !valid {
+				return i, goerrors.InvalidSyntaxError("invalid GO repeat count: GO takes a positive whole number", loc, "")
+			}
+			return i, goerrors.InvalidSyntaxError("GO batch separator inside a statement: a line holding only GO ends the batch", loc, "")
+		}
+	}
+	return -1, parseErr
+}
+
 // noteStatementEnd is called by the statement loops after a statement was
 // parsed. When no ';' follows and the next token unambiguously starts a new
 // statement, that token is marked as a statement boundary so that utility

pkg/sql/parser/statement_boundary_test.go

@@ -400,3 +400,78 @@ func TestParser_UnclosedBlocksBoundedWork(t *testing.T) {
 		})
 	}
 }
+
+// TestParser_GOLineInsideStatement checks that a line holding only GO [n]
+// ends the batch wherever it appears, as sqlcmd splits the script there: a
+// statement that would run across it is an error, not a statement with a
+// column or alias named go. GO with an invalid repeat count gets its own
+// error instead of a confusing one from the statement parser.
+func TestParser_GOLineInsideStatement(t *testing.T) {
+	tests := []struct {
+		name string
+		sql  string
+		want string // substring of the error
+	}{
+		{"select list", "SELECT a,\ngo\nFROM t", "GO batch separator inside a statement"},
+		// GO ends the first batch before it; the second one starts with FROM.
+		{"alias position", "SELECT a\nGO 2\nFROM t", "FROM"},
+		{"where clause", "UPDATE t SET a = 1 WHERE\nGO\nb = 2", "GO batch separator inside a statement"},
+		{"lower case with count", "SELECT a,\n  go 3\nFROM t", "GO batch separator inside a statement"},
+		{"zero count", "SELECT 1\nGO 0\nSELECT 2", "invalid GO repeat count"},
+		{"fractional count", "SELECT 1\nGO 1.5\nSELECT 2", "invalid GO repeat count"},
+		{"zero count first", "GO 0\nSELECT 2", "invalid GO repeat count"},
+		{"zero count last", "SELECT 1\nGO 0", "invalid GO repeat count"},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			tree, err := ParseWithDialect(tc.sql, keywords.DialectSQLServer)
+			if err == nil {
+				defer ast.ReleaseAST(tree)
+				t.Fatalf("parsed as %v, want an error", stmtShapes(tree))
+			}
+			if !strings.Contains(err.Error(), tc.want) {
+				t.Fatalf("error = %v, want it to mention %q", err, tc.want)
+			}
+		})
+	}
+
+	// Other dialects have no GO directive: go is an ordinary name there.
+	tree, err := ParseWithDialect("SELECT a,\ngo\nFROM t", keywords.DialectPostgreSQL)
+	if err != nil {
+		t.Fatalf("postgresql: %v", err)
+	}
+	ast.ReleaseAST(tree)
+
+	// GO on the same line as other text is not a separator.
+	for _, sql := range []string{"SELECT a, go FROM t", "SELECT a,\ngo, b\nFROM t"} {
+		tree, err := ParseWithDialect(sql, keywords.DialectSQLServer)
+		if err != nil {
+			t.Fatalf("%q: %v", sql, err)
+		}
+		ast.ReleaseAST(tree)
+	}
+}
+
+// TestParser_RecoveryStopsAtGOLineInsideStatement checks that the recovering
+// parser reports a statement that ran across a GO line and keeps the batch
+// boundary.
+func TestParser_RecoveryStopsAtGOLineInsideStatement(t *testing.T) {
+	tkz := tokenizer.GetTokenizer()
+	defer tokenizer.PutTokenizer(tkz)
+	tkz.SetDialect(keywords.DialectSQLServer)
+	tokens, err := tkz.TokenizeString("SELECT a,\nGO\nFROM t\nGO\nSELECT 2")
+	if err != nil {
+		t.Fatal(err)
+	}
+	p := NewParser(WithDialect("sqlserver"))
+	defer p.Release()
+	stmts, errs := p.ParseWithRecoveryFromModelTokens(tokens)
+	tree := &ast.AST{Statements: stmts}
+	defer ast.ReleaseAST(tree)
+	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "GO batch separator inside a statement") {
+		t.Fatalf("errors = %v, want a GO separator error first", errs)
+	}
+	if got := fmt.Sprint(stmtShapes(tree)); !strings.HasPrefix(got, "[BatchSeparator") || !strings.HasSuffix(got, "BatchSeparator Select]") {
+		t.Fatalf("recovery shapes = %s", got)
+	}
+}