generic setMySQLOptimizerHint

Signed-off-by: Tim Vaillancourt <[email protected]>

go/vt/sqlparser/comments.go

@@ -324,29 +324,29 @@ func (c *ParsedComments) GetMySQLSetVarValue(key string) string {
 	return ""
 }
 
-// SetMySQLMaxExecutionTime sets a query level maximum execution time using a /*+ MAX_EXECUTION_TIME() */ MySQL optimizer hint.
-func (c *ParsedComments) SetMySQLMaxExecutionTime(maxExecutionTime time.Duration) (newComments Comments) {
-	return c.SetMySQLOptimizerHint(OptimizerHintMaxExecutionTime, "", maxExecutionTime.Milliseconds())
+// SetMySQLMaxExecutionTimeValue sets the maximum execution time for a query using a /*+ MAX_EXECUTION_TIME() */ MySQL optimizer hint.
+func (c *ParsedComments) SetMySQLMaxExecutionTimeValue(maxExecutionTime time.Duration) (newComments Comments) {
+	return setMySQLOptimizerHint(c.comments, OptimizerHintMaxExecutionTime, "" /* no key */, maxExecutionTime.Milliseconds())
 }
 
 // SetMySQLSetVarValue updates or sets the value of the given variable as part of a /*+ SET_VAR() */ MySQL optimizer hint.
 func (c *ParsedComments) SetMySQLSetVarValue(key string, value string) (newComments Comments) {
-	return c.SetMySQLOptimizerHint(OptimizerHintSetVar, key, value)
+	return setMySQLOptimizerHint(c.comments, OptimizerHintSetVar, key, value)
 }
 
-// SetMySQLOptimizerHint updates or sets the value of a MySQL optimizer hint.
-func (c *ParsedComments) SetMySQLOptimizerHint(hint, key string, value interface{}) (newComments Comments) {
+// setMySQLOptimizerHint updates or sets the value of a MySQL optimizer hint.
+func setMySQLOptimizerHint(comments Comments, hint, key string, value interface{}) (newComments Comments) {
 	keyAndValue := value
 	if key != "" {
 		keyAndValue = fmt.Sprintf("%v=%v", key, value)
 	}
-	if c == nil {
+	if len(comments) == 0 {
 		// If we have no parsed comments, then we create a new one with the required optimizer hint and return it.
 		newComments = append(newComments, fmt.Sprintf("/*+ %v(%v) */", hint, keyAndValue))
 		return
 	}
 	seenFirstOhComment := false
-	for _, commentStr := range c.comments {
+	for _, commentStr := range comments {
 		// Skip all the comments that don't start with the query optimizer prefix.
 		// Also, since MySQL only parses the first comment that has the optimizer hint prefix and ignores the following ones,
 		// we skip over all the comments that come after we have seen the first comment with the optimizer hint.

go/vt/sqlparser/comments_test.go

@@ -601,7 +601,7 @@ func TestGetMySQLSetVarValue(t *testing.T) {
 	}
 }
 
-func TestSetMySQLMaxExecutionTime(t *testing.T) {
+func TestSetMySQLMaxExecutionTimeValue(t *testing.T) {
 	tests := []struct {
 		name           string
 		comments       []string
@@ -617,16 +617,16 @@ func TestSetMySQLMaxExecutionTime(t *testing.T) {
 		{
 			name:           "Add to comments",
 			comments:       []string{"/*+ SET_VAR(sort_buffer_size = 16M) */"},
-			maxExecTime:    time.Second * 30,
-			commentsWanted: []string{"/*+ SET_VAR(sort_buffer_size = 16M) MAX_EXECUTION_TIME(30000) */"},
+			maxExecTime:    time.Minute,
+			commentsWanted: []string{"/*+ SET_VAR(sort_buffer_size = 16M) MAX_EXECUTION_TIME(60000) */"},
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &ParsedComments{
 				comments: tt.comments,
 			}
-			newComments := c.SetMySQLMaxExecutionTime(tt.maxExecTime)
+			newComments := c.SetMySQLMaxExecutionTimeValue(tt.maxExecTime)
 			require.EqualValues(t, tt.commentsWanted, newComments)
 		})
 	}