From f58da1a923397ac8b1ef7a4b6b8f78d22de80b6c Mon Sep 17 00:00:00 2001
From: David Shiflet <david.shiflet@microsoft.com>
Date: Tue, 1 Aug 2023 10:58:28 -0500
Subject: [PATCH] Fix: Use proper reader for long file lines (#441)

---
 pkg/sqlcmd/sqlcmd.go      | 33 ++++++++++++++++++---------------
 pkg/sqlcmd/sqlcmd_test.go | 16 ++++++++++++++++
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/pkg/sqlcmd/sqlcmd.go b/pkg/sqlcmd/sqlcmd.go
index d7519b42..051cd6bc 100644
--- a/pkg/sqlcmd/sqlcmd.go
+++ b/pkg/sqlcmd/sqlcmd.go
@@ -39,8 +39,6 @@ var (
 	}
 )
 
-const maxLineBuffer = 2 * 1024 * 1024 // 2Mb
-
 // Console defines methods used for console input and output
 type Console interface {
 	// Readline returns the next line of input.
@@ -327,24 +325,29 @@ func (s *Sqlcmd) IncludeFile(path string, processAll bool) error {
 	b := s.batch.batchline
 	utf16bom := unicode.BOMOverride(unicode.UTF8.NewDecoder())
 	unicodeReader := transform.NewReader(f, utf16bom)
-	scanner := bufio.NewScanner(unicodeReader)
-	buf := make([]byte, maxLineBuffer)
-	scanner.Buffer(buf, maxLineBuffer)
+	scanner := bufio.NewReader(unicodeReader)
 	curLine := s.batch.read
 	echoFileLines := s.echoFileLines
+	ln := make([]byte, 0, 2*1024*1024)
 	s.batch.read = func() (string, error) {
-		if !scanner.Scan() {
-			err := scanner.Err()
-			if err == nil {
-				return "", io.EOF
-			}
-			return "", err
+		var (
+			isPrefix bool  = true
+			err      error = nil
+			line     []byte
+		)
+
+		for isPrefix && err == nil {
+			line, isPrefix, err = scanner.ReadLine()
+			ln = append(ln, line...)
 		}
-		t := scanner.Text()
-		if echoFileLines {
-			_, _ = s.GetOutput().Write([]byte(s.Prompt() + t + SqlcmdEol))
+		if err == nil && echoFileLines {
+			_, _ = s.GetOutput().Write([]byte(s.Prompt()))
+			_, _ = s.GetOutput().Write(ln)
+			_, _ = s.GetOutput().Write([]byte(SqlcmdEol))
 		}
-		return t, nil
+		t := string(ln)
+		ln = ln[:0]
+		return t, err
 	}
 	err = s.Run(false, processAll)
 	s.batch.read = curLine
diff --git a/pkg/sqlcmd/sqlcmd_test.go b/pkg/sqlcmd/sqlcmd_test.go
index 231f9b02..6af54f6a 100644
--- a/pkg/sqlcmd/sqlcmd_test.go
+++ b/pkg/sqlcmd/sqlcmd_test.go
@@ -549,6 +549,22 @@ func TestSqlCmdOutputAndError(t *testing.T) {
 	}
 }
 
+func TestVeryLongLineInFile(t *testing.T) {
+	s, buf := setupSqlCmdWithMemoryOutput(t)
+	val := strings.Repeat("a1b", (3*1024*1024)/3)
+	line := "set nocount on" + SqlcmdEol + "select('" + val + "')"
+	file, err := os.CreateTemp("", "sqlcmdlongline")
+	assert.NoError(t, err, "os.CreateTemp")
+	defer os.Remove(file.Name())
+	_, err = file.WriteString(line)
+	assert.NoError(t, err, "Unable to write temp file")
+	err = s.IncludeFile(file.Name(), true)
+	if assert.NoError(t, err, "runSqlCmd") {
+		actual := strings.TrimRight(buf.buf.String(), "\r\n")
+		assert.Equal(t, val, actual, "Query result")
+	}
+}
+
 // runSqlCmd uses lines as input for sqlcmd instead of relying on file or console input
 func runSqlCmd(t testing.TB, s *Sqlcmd, lines []string) error {
 	t.Helper()