From 11f241bb381482a0077b887633ad29876c551f88 Mon Sep 17 00:00:00 2001
From: davidshi <david.shiflet@microsoft.com>
Date: Wed, 27 Sep 2023 11:01:27 -0500
Subject: [PATCH 1/2] allow tabs in commands

---
 pkg/sqlcmd/commands_test.go | 3 ++-
 pkg/sqlcmd/parse.go         | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/pkg/sqlcmd/commands_test.go b/pkg/sqlcmd/commands_test.go
index 1119df93..94caeb2e 100644
--- a/pkg/sqlcmd/commands_test.go
+++ b/pkg/sqlcmd/commands_test.go
@@ -43,10 +43,11 @@ func TestCommandParsing(t *testing.T) {
 		{`:Setvar A1 "some value" `, "SETVAR", []string{`A1 "some value" `}},
 		{` :Listvar`, "LISTVAR", []string{""}},
 		{`:EXIT (select 100 as count)`, "EXIT", []string{" (select 100 as count)"}},
+		{"\t:EXIT (select 100 as count)", "EXIT", []string{" (select 100 as count)"}},
 		{`:EXIT ( )`, "EXIT", []string{" ( )"}},
 		{`EXIT `, "EXIT", []string{" "}},
 		{`:Connect someserver -U someuser`, "CONNECT", []string{"someserver -U someuser"}},
-		{`:r c:\$(var)\file.sql`, "READFILE", []string{`c:\$(var)\file.sql`}},
+		{":r\tc:\\$(var)\\file.sql", "READFILE", []string{`c:\$(var)\file.sql`}},
 		{`:!! notepad`, "EXEC", []string{" notepad"}},
 		{`:!!notepad`, "EXEC", []string{"notepad"}},
 		{` !! dir c:\`, "EXEC", []string{` dir c:\`}},
diff --git a/pkg/sqlcmd/parse.go b/pkg/sqlcmd/parse.go
index c1c10ef4..11dd14e1 100644
--- a/pkg/sqlcmd/parse.go
+++ b/pkg/sqlcmd/parse.go
@@ -34,7 +34,7 @@ func readMultilineComment(r []rune, i, end int) (int, bool) {
 func readCommand(c Commands, r []rune, i, end int) (*Command, []string, int) {
 	for ; i < end; i++ {
 		next := grab(r, i, end)
-		if next == 0 || unicode.IsControl(next) {
+		if next == 0 || (unicode.IsControl(next) && next != '\t') {
 			break
 		}
 	}

From d67140d5ed77bc78272fcefb42ef5ebf699066cd Mon Sep 17 00:00:00 2001
From: davidshi <david.shiflet@microsoft.com>
Date: Fri, 29 Sep 2023 14:06:38 -0500
Subject: [PATCH 2/2] Fix variable substitution

---
 pkg/sqlcmd/sqlcmd.go      | 8 +++++---
 pkg/sqlcmd/sqlcmd_test.go | 1 +
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/pkg/sqlcmd/sqlcmd.go b/pkg/sqlcmd/sqlcmd.go
index a4e37c07..3293253b 100644
--- a/pkg/sqlcmd/sqlcmd.go
+++ b/pkg/sqlcmd/sqlcmd.go
@@ -389,6 +389,8 @@ func (s *Sqlcmd) getRunnableQuery(q string) string {
 	}
 	b := new(strings.Builder)
 	b.Grow(len(q))
+	// The varmap index is rune based not byte based
+	r := []rune(q)
 	keys := make([]int, 0, len(s.batch.varmap))
 	for k := range s.batch.varmap {
 		keys = append(keys, k)
@@ -396,7 +398,7 @@ func (s *Sqlcmd) getRunnableQuery(q string) string {
 	sort.Ints(keys)
 	last := 0
 	for _, i := range keys {
-		b.WriteString(q[last:i])
+		b.WriteString(string(r[last:i]))
 		v := s.batch.varmap[i]
 		if val, ok := s.resolveVariable(v); ok {
 			b.WriteString(val)
@@ -404,9 +406,9 @@ func (s *Sqlcmd) getRunnableQuery(q string) string {
 			_, _ = fmt.Fprintf(s.GetError(), "'%s' scripting variable not defined.%s", v, SqlcmdEol)
 			b.WriteString(fmt.Sprintf("$(%s)", v))
 		}
-		last = i + len(v) + 3
+		last = i + len([]rune(v)) + 3
 	}
-	b.WriteString(q[last:])
+	b.WriteString(string(r[last:]))
 	return b.String()
 }
 
diff --git a/pkg/sqlcmd/sqlcmd_test.go b/pkg/sqlcmd/sqlcmd_test.go
index fbc5c915..3e4729b5 100644
--- a/pkg/sqlcmd/sqlcmd_test.go
+++ b/pkg/sqlcmd/sqlcmd_test.go
@@ -246,6 +246,7 @@ func TestGetRunnableQuery(t *testing.T) {
 		{"$ (var2)", "$ (var2)"},
 		{"select '$(VAR1) $(VAR2)' as  c", "select 'v1 variable2' as  c"},
 		{" $(VAR1) ' $(VAR2) ' as  $(VAR1)", " v1 ' variable2 ' as  v1"},
+		{"í $(VAR1)", "í v1"},
 	}
 	s := New(nil, "", v)
 	for _, test := range tests {