From 0221d93de85506dbb3564603666f60708ce6804c Mon Sep 17 00:00:00 2001
From: David Shiflet <david.shiflet@microsoft.com>
Date: Fri, 23 Aug 2024 09:38:09 -0500
Subject: [PATCH] omit comments from commands (#548)

---
 pkg/sqlcmd/commands.go      | 43 +++++++++++++++++++++++++++-
 pkg/sqlcmd/commands_test.go | 56 +++++++++++++++++++++++++++++++++++--
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/pkg/sqlcmd/commands.go b/pkg/sqlcmd/commands.go
index 31c2d9c0..9bf0b9d9 100644
--- a/pkg/sqlcmd/commands.go
+++ b/pkg/sqlcmd/commands.go
@@ -134,12 +134,53 @@ func (c Commands) matchCommand(line string) (*Command, []string) {
 	for _, cmd := range c {
 		matchedCommand := cmd.regex.FindStringSubmatch(line)
 		if matchedCommand != nil {
-			return cmd, matchedCommand[1:]
+			return cmd, removeComments(matchedCommand[1:])
 		}
 	}
 	return nil, nil
 }
 
+func removeComments(args []string) []string {
+	var pos int
+	quote := false
+	for i := range args {
+		pos, quote = commentStart([]rune(args[i]), quote)
+		if pos > -1 {
+			out := make([]string, i+1)
+			if i > 0 {
+				copy(out, args[:i])
+			}
+			out[i] = args[i][:pos]
+			return out
+		}
+	}
+	return args
+}
+
+func commentStart(arg []rune, quote bool) (int, bool) {
+	var i int
+	space := true
+	for ; i < len(arg); i++ {
+		c, next := arg[i], grab(arg, i+1, len(arg))
+		switch {
+		case quote && c == '"' && next != '"':
+			quote = false
+		case quote && c == '"' && next == '"':
+			i++
+		case c == '\t' || c == ' ':
+			space = true
+		// Note we assume none of the regexes would split arguments on non-whitespace boundaries such that "text -- comment" would get split into "text -" and "- comment"
+		case !quote && space && c == '-' && next == '-':
+			return i, false
+		case !quote && c == '"':
+			quote = true
+		default:
+			space = false
+		}
+	}
+	return -1, quote
+}
+
 func warnDisabled(s *Sqlcmd, args []string, line uint) error {
 	s.WriteError(s.GetError(), ErrCommandsDisabled)
 	return nil
diff --git a/pkg/sqlcmd/commands_test.go b/pkg/sqlcmd/commands_test.go
index 7a9e8f3f..7d2acb4e 100644
--- a/pkg/sqlcmd/commands_test.go
+++ b/pkg/sqlcmd/commands_test.go
@@ -59,8 +59,18 @@ func TestCommandParsing(t *testing.T) {
 		cmd, args := c.matchCommand(test.line)
 		if test.cmd != "" {
 			if assert.NotNil(t, cmd, "No command found for `%s`", test.line) {
-				assert.Equal(t, test.cmd, cmd.name, "Incorrect command for `%s`", test.line)
-				assert.Equal(t, test.args, args, "Incorrect arguments for `%s`", test.line)
+				assert.Equalf(t, test.cmd, cmd.name, "Incorrect command for `%s`", test.line)
+				assert.Equalf(t, test.args, args, "Incorrect arguments for `%s`", test.line)
+				line := test.line + " -- comment"
+				cmd, args = c.matchCommand(line)
+				if assert.NotNil(t, cmd, "No command found for `%s`", line) {
+					assert.Equalf(t, test.cmd, cmd.name, "Incorrect command for `%s`", line)
+					assert.Equalf(t, len(test.args), len(args), "Incorrect argument count for `%s`.", line)
+					for _, a := range args {
+						assert.NotContains(t, a, "--", "comment marker should be omitted")
+						assert.NotContains(t, a, "comment", "comment should e omitted")
+					}
+				}
 			}
 		} else {
 			assert.Nil(t, cmd, "Unexpected match for %s", test.line)
@@ -68,6 +78,48 @@ func TestCommandParsing(t *testing.T) {
 	}
 }
 
+func TestRemoveComments(t *testing.T) {
+	type testData struct {
+		args   []string
+		result []string
+	}
+	tests := []testData{
+		{[]string{"-- comment"}, []string{""}},
+		{[]string{"filename -- comment"}, []string{"filename "}},
+		{[]string{`"file""name"`, `-- comment`}, []string{`"file""name"`, ""}},
+		{[]string{`"file""name"--comment`}, []string{`"file""name"--comment`}},
+	}
+	for _, test := range tests {
+		actual := removeComments(test.args)
+		assert.Equal(t, test.result, actual, "Comments not removed properly")
+	}
+}
+
+func TestCommentStart(t *testing.T) {
+	type testData struct {
+		arg      string
+		quoteIn  bool
+		quoteOut bool
+		pos      int
+	}
+	tests := []testData{
+		{"nospace-- comment", false, false, -1},
+		{"-- comment", false, false, 0},
+		{"-- comment", true, true, -1},
+		{`" ""quoted""`, false, true, -1},
+		{`"-- ""quoted""`, false, true, -1},
+		{`"-- ""quoted"" " -- comment`, false, false, 17},
+		{`"-- ""quoted"" " -- comment`, true, false, 1},
+	}
+	for _, test := range tests {
+		t.Run(test.arg, func(t *testing.T) {
+			i, q := commentStart([]rune(test.arg), test.quoteIn)
+			assert.Equal(t, test.quoteOut, q, "Wrong quote")
+			assert.Equal(t, test.pos, i, "Wrong position")
+		})
+	}
+}
+
 func TestCustomBatchSeparator(t *testing.T) {
 	c := newCommands()
 	err := c.SetBatchTerminator("me!")