fix: insert on duplicate key update missing BindVars (#14728)

Co-authored-by: wlx5575 <[email protected]>
vitessio · Dec 12, 2023 · 56aba1f · 56aba1f
1 parent 21b490a
commit 56aba1f
Show file tree

Hide file tree

Showing 10 changed files with 327 additions and 102 deletions.
diff --git a/go/test/endtoend/vtgate/misc_test.go b/go/test/endtoend/vtgate/misc_test.go
@@ -31,6 +31,16 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestInsertOnDuplicateKey(t *testing.T) {
+	conn, closer := start(t)
+	defer closer()
+
+	utils.Exec(t, conn, "insert into t11(id, sharding_key, col1, col2, col3) values(1, 2, 'a', 1, 2)")
+	utils.Exec(t, conn, "insert into t11(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 t11 where id=10", "[[INT64(10) INT64(2)]]")
+
+}
+
 func TestInsertNeg(t *testing.T) {
 	conn, closer := start(t)
 	defer closer()

diff --git a/go/test/endtoend/vtgate/schema.sql b/go/test/endtoend/vtgate/schema.sql
@@ -155,3 +155,13 @@ create table t10_id_to_keyspace_id_idx
     keyspace_id varbinary(10),
     primary key (id)
 ) Engine = InnoDB;
+
+create table t11
+(
+    id           bigint,
+    sharding_key bigint,
+    col1 varchar(50),
+    col2 int,
+    col3 int,
+    primary key (id)
+) Engine = InnoDB;
diff --git a/go/test/endtoend/vtgate/vschema.json b/go/test/endtoend/vtgate/vschema.json
@@ -1,11 +1,10 @@
-
 {
   "sharded": true,
   "vindexes": {
-    "unicode_loose_xxhash" : {
+    "unicode_loose_xxhash": {
       "type": "unicode_loose_xxhash"
     },
-    "unicode_loose_md5" : {
+    "unicode_loose_md5": {
       "type": "unicode_loose_md5"
     },
     "hash": {
@@ -159,7 +158,10 @@
           "name": "hash"
         },
         {
-          "columns": ["id2", "id1"],
+          "columns": [
+            "id2",
+            "id1"
+          ],
           "name": "t4_id2_vdx"
         }
       ]
@@ -179,7 +181,10 @@
           "name": "hash"
         },
         {
-          "columns": ["id2", "id1"],
+          "columns": [
+            "id2",
+            "id1"
+          ],
           "name": "t6_id2_vdx"
         }
       ]
@@ -301,6 +306,14 @@
           "name": "hash"
         }
       ]
+    },
+    "t11": {
+      "column_vindexes": [
+        {
+          "column": "sharding_key",
+          "name": "hash"
+        }
+      ]
     }
   }
 }
diff --git a/go/vt/vtgate/engine/cached_size.go b/go/vt/vtgate/engine/cached_size.go
diff --git a/go/vt/vtgate/engine/insert.go b/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,

diff --git a/go/vt/vtgate/engine/insert_common.go b/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

diff --git a/go/vt/vtgate/engine/insert_select.go b/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,