diff --git a/pkg/sqlcmd/batch_test.go b/pkg/sqlcmd/batch_test.go index e3175245..3748d2b0 100644 --- a/pkg/sqlcmd/batch_test.go +++ b/pkg/sqlcmd/batch_test.go @@ -40,6 +40,7 @@ func TestBatchNext(t *testing.T) { {"exit_1", []string{"exit_1"}, nil, "-"}, {"print ('me\n!!me')", []string{"print ('me" + SqlcmdEol + "!!me')"}, nil, "-"}, {"/*comment\n!!me*/", []string{"/*comment" + SqlcmdEol + "!!me*/"}, nil, "-"}, + {"select 1\nGO\t\t", []string{"select 1"}, []string{"GO"}, "-"}, } for _, test := range tests { b := NewBatch(sp(test.s, "\n"), newCommands()) diff --git a/pkg/sqlcmd/commands.go b/pkg/sqlcmd/commands.go index 66dd1dba..6826edba 100644 --- a/pkg/sqlcmd/commands.go +++ b/pkg/sqlcmd/commands.go @@ -193,7 +193,7 @@ func errorDisabled(s *Sqlcmd, args []string, line uint) error { } func batchTerminatorRegex(terminator string) string { - return fmt.Sprintf(`(?im)^[\t ]*?%s(?:[ ]+(.*$)|$)`, regexp.QuoteMeta(terminator)) + return fmt.Sprintf(`(?im)^[\t ]*?%s(?:[\t ]+(.*$)|$)`, regexp.QuoteMeta(terminator)) } // SetBatchTerminator attempts to set the batch terminator to the given value diff --git a/pkg/sqlcmd/commands_test.go b/pkg/sqlcmd/commands_test.go index 76c509a8..6f6d8689 100644 --- a/pkg/sqlcmd/commands_test.go +++ b/pkg/sqlcmd/commands_test.go @@ -54,6 +54,10 @@ func TestCommandParsing(t *testing.T) { {`:XML ON `, "XML", []string{`ON `}}, {`:RESET`, "RESET", []string{""}}, {`RESET`, "RESET", []string{""}}, + {"GO", "GO", []string{""}}, + {"GO\t\t", "GO", []string{""}}, + {"go\t5", "GO", []string{"5"}}, + {" \tGO \t 5", "GO", []string{"5"}}, } for _, test := range commands { @@ -130,6 +134,10 @@ func TestCustomBatchSeparator(t *testing.T) { assert.Equal(t, "GO", cmd.name, "command name") assert.Equal(t, "5", strings.TrimSpace(args[0]), "go argument") } + cmd, args = c.matchCommand("me!\t5") + if assert.NotNil(t, cmd, "matchCommand didn't find custom batch separator followed by a tab") { + assert.Equal(t, "5", strings.TrimSpace(args[0]), "go argument") + } } }