fix: insert on duplicate key update missing BindVars

Signed-off-by: wlx5575 <[email protected]>

go/test/endtoend/vtgate/misc_test.go

@@ -27,6 +27,16 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestInsertOnDuplicateKey(t *testing.T) {
+	conn, closer := start(t)
+	defer closer()
+
+	utils.Exec(t, conn, "insert into t10(id, sharding_key, col1, col2, col3) values(1, 2, 'a', 1, 2)")
+	utils.Exec(t, conn, "insert into t10(id, sharding_key, col1, col2, col3) values(1, 2, 'a', 1, 2) on duplicate key update id=10;")
+	utils.AssertMatches(t, conn, "select id, sharding_key from t10 where id=10", "[[INT64(10) INT64(2)] ")
+
+}
+
 func TestInsertNeg(t *testing.T) {
 	conn, closer := start(t)
 	defer closer()

go/vt/vtgate/engine/insert.go

@@ -78,7 +78,7 @@ func newInsert(
 	table *vindexes.Table,
 	prefix string,
 	mid sqlparser.Values,
-	suffix string,
+	suffix sqlparser.OnDup,
 ) *Insert {
 	ins := &Insert{
 		InsertCommon: InsertCommon{
@@ -265,25 +265,30 @@ func (ins *Insert) getInsertShardedQueries(
 		for _, indexValue := range indexesPerRss[i] {
 			index, _ := strconv.ParseInt(string(indexValue.Value), 0, 64)
 			if keyspaceIDs[index] != nil {
+				walkFunc := func(node sqlparser.SQLNode) (kontinue bool, err error) {
+					if arg, ok := node.(*sqlparser.Argument); ok {
+						bv, exists := bindVars[arg.Name]
+						if !exists {
+							return false, vterrors.VT03026(arg.Name)
+						}
+						shardBindVars[arg.Name] = bv
+					}
+					return true, nil
+				}
 				mids = append(mids, sqlparser.String(ins.Mid[index]))
 				for _, expr := range ins.Mid[index] {
-					err = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
-						if arg, ok := node.(*sqlparser.Argument); ok {
-							bv, exists := bindVars[arg.Name]
-							if !exists {
-								return false, vterrors.VT03026(arg.Name)
-							}
-							shardBindVars[arg.Name] = bv
-						}
-						return true, nil
-					}, expr, nil)
+					err = sqlparser.Walk(walkFunc, expr, nil)
 					if err != nil {
 						return nil, nil, err
 					}
 				}
+				err = sqlparser.Walk(walkFunc, ins.Suffix, nil)
+				if err != nil {
+					return nil, nil, err
+				}
 			}
 		}
-		rewritten := ins.Prefix + strings.Join(mids, ",") + ins.Suffix
+		rewritten := ins.Prefix + strings.Join(mids, ",") + sqlparser.String(ins.Suffix)
 		queries[i] = &querypb.BoundQuery{
 			Sql:           rewritten,
 			BindVariables: shardBindVars,

go/vt/vtgate/engine/insert_common.go

@@ -23,6 +23,8 @@ import (
 	"strconv"
 	"strings"
 
+	"vitess.io/vitess/go/vt/sqlparser"
+
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -73,7 +75,7 @@ type (
 
 		// Prefix, Suffix are for sharded insert plans.
 		Prefix string
-		Suffix string
+		Suffix sqlparser.OnDup
 
 		// Insert needs tx handling
 		txNeeded

go/vt/vtgate/engine/insert_select.go

@@ -56,7 +56,7 @@ func newInsertSelect(
 	keyspace *vindexes.Keyspace,
 	table *vindexes.Table,
 	prefix string,
-	suffix string,
+	suffix sqlparser.OnDup,
 	vv [][]int,
 	input Primitive,
 ) *InsertSelect {
@@ -172,7 +172,7 @@ func (ins *InsertSelect) getInsertUnshardedQuery(rows []sqltypes.Row, bindVars m
 		}
 		mids = append(mids, row)
 	}
-	return ins.Prefix + sqlparser.String(mids) + ins.Suffix
+	return ins.Prefix + sqlparser.String(mids) + sqlparser.String(ins.Suffix)
 }
 
 func (ins *InsertSelect) insertIntoShardedTable(
@@ -270,7 +270,7 @@ func (ins *InsertSelect) getInsertShardedQueries(
 				mids = append(mids, row)
 			}
 		}
-		rewritten := ins.Prefix + sqlparser.String(mids) + ins.Suffix
+		rewritten := ins.Prefix + sqlparser.String(mids) + sqlparser.String(ins.Suffix)
 		queries[i] = &querypb.BoundQuery{
 			Sql:           rewritten,
 			BindVariables: bvs,