Fix issues with special characters in files (#463)

microsoft · Oct 10, 2023 · a62432f · a62432f
1 parent 53b6fe7
commit a62432f
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 5 deletions.
diff --git 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
@@ -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
 		}
 	}

diff --git a/pkg/sqlcmd/sqlcmd.go b/pkg/sqlcmd/sqlcmd.go
@@ -389,24 +389,26 @@ 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)
 	}
 	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)
 		} else {
 			_, _ = 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
@@ -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 {