From c15cebdaf73bb1a7b010ceb32830919c54f0d5ee Mon Sep 17 00:00:00 2001
From: davidshi <david.shiflet@microsoft.com>
Date: Mon, 31 Jul 2023 17:50:06 -0500
Subject: [PATCH 1/2] Fix: Use proper reader for long file lines

---
 pkg/sqlcmd/sqlcmd.go      | 28 +++++++++++++++-------------
 pkg/sqlcmd/sqlcmd_test.go | 16 ++++++++++++++++
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/pkg/sqlcmd/sqlcmd.go b/pkg/sqlcmd/sqlcmd.go
index d7519b42..e6f9f69a 100644
--- a/pkg/sqlcmd/sqlcmd.go
+++ b/pkg/sqlcmd/sqlcmd.go
@@ -327,24 +327,26 @@ 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
 	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, ln []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
+		return string(ln), 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()

From fef7332eb3ffd9708357b568cafc02daa3c9e815 Mon Sep 17 00:00:00 2001
From: davidshi <david.shiflet@microsoft.com>
Date: Mon, 31 Jul 2023 20:29:44 -0500
Subject: [PATCH 2/2] init line buffer to 2MB

---
 pkg/sqlcmd/sqlcmd.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/pkg/sqlcmd/sqlcmd.go b/pkg/sqlcmd/sqlcmd.go
index e6f9f69a..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.
@@ -330,11 +328,12 @@ func (s *Sqlcmd) IncludeFile(path string, processAll bool) error {
 	scanner := bufio.NewReader(unicodeReader)
 	curLine := s.batch.read
 	echoFileLines := s.echoFileLines
+	ln := make([]byte, 0, 2*1024*1024)
 	s.batch.read = func() (string, error) {
 		var (
 			isPrefix bool  = true
 			err      error = nil
-			line, ln []byte
+			line     []byte
 		)
 
 		for isPrefix && err == nil {
@@ -346,7 +345,9 @@ func (s *Sqlcmd) IncludeFile(path string, processAll bool) error {
 			_, _ = s.GetOutput().Write(ln)
 			_, _ = s.GetOutput().Write([]byte(SqlcmdEol))
 		}
-		return string(ln), err
+		t := string(ln)
+		ln = ln[:0]
+		return t, err
 	}
 	err = s.Run(false, processAll)
 	s.batch.read = curLine