Merge pull request #3848 from tinyspeck/on-dup-shard-key

Insert on duplicate key query support enhancement

data/test/vtgate/dml_cases.txt

@@ -464,6 +464,73 @@
   }
 }
 
+# sharded upsert with sharding key set to vindex column
+"insert into music(user_id, id) values(1, 2) on duplicate key update user_id = values(user_id)"
+{
+  "Original": "insert into music(user_id, id) values(1, 2) on duplicate key update user_id = values(user_id)",
+  "Instructions": {
+    "Opcode": "InsertShardedIgnore",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "insert into music(user_id, id) values (:_user_id0, :_id0) on duplicate key update user_id = values(user_id)",
+    "Values": [
+      [
+        [
+          1
+        ]
+      ],
+      [
+        [
+          2
+        ]
+      ]
+    ],
+    "Table": "music",
+    "Prefix": "insert into music(user_id, id) values ",
+    "Mid": [
+      "(:_user_id0, :_id0)"
+    ],
+    "Suffix": " on duplicate key update user_id = values(user_id)"
+  }
+}
+
+# sharded bulk upsert with sharding key set to vindex column 
+"insert into music(user_id, id) values (1, 2), (3,4) on duplicate key update user_id = values(user_id)"
+{
+  "Original": "insert into music(user_id, id) values (1, 2), (3,4) on duplicate key update user_id = values(user_id)",
+  "Instructions": {
+    "Opcode": "InsertShardedIgnore",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "insert into music(user_id, id) values (:_user_id0, :_id0), (:_user_id1, :_id1) on duplicate key update user_id = values(user_id)",
+    "Values": [
+      [
+        [
+          1,
+          3
+        ]
+      ],
+      [
+        [
+          2,
+          4
+        ]
+      ]
+    ],
+    "Table": "music",
+    "Prefix": "insert into music(user_id, id) values ",
+    "Mid": [
+      "(:_user_id0, :_id0)",
+      "(:_user_id1, :_id1)"
+    ],
+    "Suffix": " on duplicate key update user_id = values(user_id)"
+  }
+}
+
 # insert unsharded with select
 "insert into unsharded select id from unsharded_auto"
 {

data/test/vtgate/unsupported_cases.txt

@@ -395,6 +395,10 @@
 "insert into user(id) values(1) on duplicate key update id = 3"
 "unsupported: DML cannot change vindex column"
 
+# sharded upsert can't change vindex using values function
+"insert into music(user_id, id) values(1, 2) on duplicate key update user_id = values(id)"
+"unsupported: DML cannot change vindex column"
+
 # sharded insert from select
 "insert into user(id) select 1 from dual"
 "unsupported: insert into select"

go/vt/vtgate/planbuilder/insert.go

@@ -235,7 +235,14 @@ func isVindexChanging(setClauses sqlparser.UpdateExprs, colVindexes []*vindexes.
 		for _, vcol := range colVindexes {
 			for _, col := range vcol.Columns {
 				if col.Equal(assignment.Name.Name) {
-					return true
+					valueExpr, isValuesFuncExpr := assignment.Expr.(*sqlparser.ValuesFuncExpr)
+					if !isValuesFuncExpr {
+						return true
+					}
+					// update on duplicate key is changing the vindex column, not supported.
+					if !valueExpr.Name.Equal(assignment.Name.Name) {
+						return true
+					}
 				}
 			}
 		}

go/vt/vtgate/planbuilder/update.go

@@ -103,7 +103,7 @@ func buildChangedVindexesValues(eupd *engine.Update, update *sqlparser.Update, c
 					return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "column has duplicate set values: '%v'", assignment.Name.Name)
 				}
 				found = true
-				pv, err := extractValueFromUpdate(assignment, vcol)
+				pv, err := extractValueFromUpdate(assignment)
 				if err != nil {
 					return nil, err
 				}
@@ -155,7 +155,7 @@ func generateUpdateSubquery(upd *sqlparser.Update, table *vindexes.Table) string
 // extractValueFromUpdate given an UpdateExpr attempts to extracts the Value
 // it's holding. At the moment it only supports: StrVal, HexVal, IntVal, ValArg.
 // If a complex expression is provided (e.g set name = name + 1), the update will be rejected.
-func extractValueFromUpdate(upd *sqlparser.UpdateExpr, col sqlparser.ColIdent) (pv sqltypes.PlanValue, err error) {
+func extractValueFromUpdate(upd *sqlparser.UpdateExpr) (pv sqltypes.PlanValue, err error) {
 	if !sqlparser.IsValue(upd.Expr) {
 		err := vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: Only values are supported. Invalid update on column: %v", upd.Name.Name)
 		return sqltypes.PlanValue{}, err