allow _ and - in variable names (#43)

Author: shueybubbles

pkg/sqlcmd/batch_test.go

@@ -175,6 +175,7 @@ func TestReadStringVarmap(t *testing.T) {
 	}
 	tests := []mapTest{
 		{`'var $(var1) var2 $(var2)'`, map[int]string{5: "var1", 18: "var2"}},
+		{`'var $(va_1) var2 $(va-2)'`, map[int]string{5: "va_1", 18: "va-2"}},
 	}
 	for _, test := range tests {
 		b := NewBatch(nil, newCommands())

pkg/sqlcmd/parse.go

@@ -4,6 +4,7 @@
 package sqlcmd
 
 import (
+	"strings"
 	"unicode"
 )
 
@@ -63,7 +64,7 @@ func readVariableReference(r []rune, i int, end int) (int, bool) {
 		if r[i] == ')' {
 			return i, true
 		}
-		if (r[i] >= 'a' && r[i] <= 'z') || (r[i] >= 'A' && r[i] <= 'Z') || (r[i] >= '0' && r[i] <= '9') {
+		if (r[i] >= 'a' && r[i] <= 'z') || (r[i] >= 'A' && r[i] <= 'Z') || (r[i] >= '0' && r[i] <= '9') || strings.ContainsRune(validVariableRunes, r[i]) {
 			continue
 		}
 		break

pkg/sqlcmd/variables.go

@@ -252,14 +252,14 @@ func (variables *Variables) Setvar(name, value string) error {
 	return nil
 }
 
+const validVariableRunes = "_-"
+
 // ValidIdentifier determines if a given string can be used as a variable name
-// The native sqlcmd allowed some punctuation characters as part of a variable name
-// but this version will not.
 func ValidIdentifier(name string) error {
 
 	first := true
 	for _, c := range name {
-		if !unicode.IsLetter(c) && (first || !unicode.IsDigit(c)) {
+		if !unicode.IsLetter(c) && (first || (!unicode.IsDigit(c) && !strings.ContainsRune(validVariableRunes, c))) {
 			return fmt.Errorf("Invalid variable identifier %s", name)
 		}
 		first = false

pkg/sqlcmd/variables_test.go

@@ -103,6 +103,7 @@ func TestValidIdentifier(t *testing.T) {
 		{"1A", false},
 		{"A1", true},
 		{"A+", false},
+		{"A-_b", true},
 	}
 	for _, tst := range tests {
 		err := ValidIdentifier(tst.raw)