Adds support to scattered updates

Signed-off-by: Rafael Chacon <rafael@slack-corp.com>

Author: Rafael Chacon

data/test/vtgate/dml_cases.txt

@@ -1287,6 +1287,117 @@
   }
 }
 
+# update with no where clause
+"update user_extra set val = 1"
+{
+  "Original": "update user_extra set val = 1",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1",
+    "Table": "user_extra"
+  }
+}
+
+
+# update with non-comparison expr
+"update user_extra set val = 1 where id between 1 and 2"
+{
+  "Original": "update user_extra set val = 1 where id between 1 and 2",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where id between 1 and 2",
+    "Table": "user_extra"
+  }
+}
+
+# update with primary id through IN clause
+"update user_extra set val = 1 where user_id in (1, 2)"
+{
+  "Original": "update user_extra set val = 1 where user_id in (1, 2)",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where user_id in (1, 2)",
+    "Table": "user_extra"
+  }
+}
+
+
+# update with non-unique key
+"update user_extra set val = 1 where name = 'foo'"
+{
+  "Original": "update user_extra set val = 1 where name = 'foo'",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where name = 'foo'",
+    "Table": "user_extra"
+  }
+}
+
+# update with no index match
+"update user_extra set val = 1 where user_id = 1"
+{
+  "Original": "update user_extra set val = 1 where user_id = 1",
+  "Instructions": {
+    "Opcode": "UpdateEqual",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where user_id = 1",
+    "Vindex": "user_index",
+    "Values": [
+      1
+    ],
+    "Table": "user_extra"
+  }
+}
+
+# update by lookup with IN clause
+"update user_extra set val = 1 where id in (1, 2)"
+{
+  "Original": "update user_extra set val = 1 where id in (1, 2)",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where id in (1, 2)",
+    "Table": "user_extra"
+  }
+}
+
+# update with where clause with parens
+"update user_extra set val = 1 where (name = 'foo' or id = 1)"
+{
+  "Original": "update user_extra set val = 1 where (name = 'foo' or id = 1)",
+  "Instructions": {
+    "Opcode": "UpdateScatter",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "Query": "update user_extra set val = 1 where (name = 'foo' or id = 1)",
+    "Table": "user_extra"
+  }
+}
+
 # delete from with no where clause
 "delete from user_extra"
 {

data/test/vtgate/unsupported_cases.txt

@@ -291,33 +291,9 @@
 "delete from unsharded where col = (select id from unsharded join user on unsharded.id = user.id)"
 "unsupported: sharded subqueries in DML"
 
-# update with no where clause
-"update user set val = 1"
-"unsupported: multi-shard where clause in DML"
-
-# update with non-comparison expr
-"update user set val = 1 where id between 1 and 2"
-"unsupported: multi-shard where clause in DML"
-
-# update with primary id through IN clause
-"update user set val = 1 where id in (1, 2)"
-"unsupported: multi-shard where clause in DML"
-
-# update with non-unique key
-"update user set val = 1 where name = 'foo'"
-"unsupported: multi-shard where clause in DML"
-
-# update with no index match
-"update user set val = 1 where user_id = 1"
-"unsupported: multi-shard where clause in DML"
-
-# update by lookup with IN clause
-"update music set val = 1 where id in (1, 2)"
-"unsupported: multi-shard where clause in DML"
-
-# update with where clause with parens
-"update user set val = 1 where (name = 'foo' or id = 1)"
-"unsupported: multi-shard where clause in DML"
+# scatter update with limit clause
+"update user_extra set val = 1 where (name = 'foo' or id = 1) limit 1"
+"unsupported: multi shard update with limit"
 
 # delete with multi-table targets
 "delete music from music where id = 1"
@@ -327,7 +303,7 @@
 "delete user from user join user_extra on user.id = user_extra.id where user.name = 'foo'"
 "unsupported: multi-table delete statement in sharded keyspace"
 
-# delete from table with no where clause and owned lookup vindex
+# scatter delete with owned lookup vindex
 "delete from user"
 "unsupported: multi shard delete on a table with owned lookup vindexes"
 
@@ -504,7 +480,3 @@
 
 "select func(keyspace_id) from user_index where id = :id"
 "unsupported: expression on results of a vindex function"
-
-# update with unsupported multi-shard where clause with parens
-"update user set val = 1 where (name = 'foo' or id = 1)"
-"unsupported: multi-shard where clause in DML"

go/vt/vtgate/engine/update.go

@@ -106,11 +106,15 @@ const (
 	// to a single shard: Requires: A Vindex, and
 	// a single Value.
 	UpdateEqual
+	// UpdateScatter is for routing a scattered
+	// update statement.
+	UpdateScatter
 )
 
 var updName = map[UpdateOpcode]string{
 	UpdateUnsharded: "UpdateUnsharded",
 	UpdateEqual:     "UpdateEqual",
+	UpdateScatter:   "UpdateScatter",
 }
 
 // MarshalJSON serializes the UpdateOpcode as a JSON string.
@@ -126,6 +130,8 @@ func (upd *Update) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVar
 		return upd.execUpdateUnsharded(vcursor, bindVars)
 	case UpdateEqual:
 		return upd.execUpdateEqual(vcursor, bindVars)
+	case UpdateScatter:
+		return upd.execUpdateByDestination(vcursor, bindVars, key.DestinationAllShards{})
 	default:
 		// Unreachable.
 		return nil, fmt.Errorf("unsupported opcode: %v", upd)
@@ -218,3 +224,21 @@ func (upd *Update) updateVindexEntries(vcursor VCursor, query string, bindVars m
 	}
 	return nil
 }
+
+func (upd *Update) execUpdateByDestination(vcursor VCursor, bindVars map[string]*querypb.BindVariable, dest key.Destination) (*sqltypes.Result, error) {
+	rss, _, err := vcursor.ResolveDestinations(upd.Keyspace.Name, nil, []key.Destination{dest})
+	if err != nil {
+		return nil, vterrors.Wrap(err, "execDeleteScatter")
+	}
+
+	queries := make([]*querypb.BoundQuery, len(rss))
+	sql := sqlannotation.AnnotateIfDML(upd.Query, nil)
+	for i := range rss {
+		queries[i] = &querypb.BoundQuery{
+			Sql:           sql,
+			BindVariables: bindVars,
+		}
+	}
+	autocommit := (len(rss) == 1) && vcursor.AutocommitApproval()
+	return vcursor.ExecuteMultiShard(rss, queries, true /* isDML */, autocommit)
+}

go/vt/vtgate/executor_dml_test.go

@@ -464,6 +464,25 @@ func TestDeleteEqual(t *testing.T) {
 	}
 }
 
+func TestUpdateScatter(t *testing.T) {
+	executor, sbc1, sbc2, _ := createExecutorEnv()
+	_, err := executorExec(executor, "update user_extra set col = 2", nil)
+	if err != nil {
+		t.Error(err)
+	}
+	// Queries get annotatted.
+	wantQueries := []*querypb.BoundQuery{{
+		Sql:           "update user_extra set col = 2/* vtgate:: filtered_replication_unfriendly */",
+		BindVariables: map[string]*querypb.BindVariable{},
+	}}
+	if !reflect.DeepEqual(sbc1.Queries, wantQueries) {
+		t.Errorf("sbc.Queries:\n%+v, want\n%+v\n", sbc1.Queries, wantQueries)
+	}
+	if !reflect.DeepEqual(sbc2.Queries, wantQueries) {
+		t.Errorf("sbc.Queries:\n%+v, want\n%+v\n", sbc2.Queries, wantQueries)
+	}
+}
+
 func TestDeleteScatter(t *testing.T) {
 	executor, sbc1, sbc2, _ := createExecutorEnv()
 	_, err := executorExec(executor, "delete from user_extra", nil)

go/vt/vtgate/planbuilder/update.go

@@ -73,9 +73,19 @@ func buildUpdatePlan(upd *sqlparser.Update, vschema ContextVSchema) (*engine.Upd
 	var err error
 	eupd.Vindex, eupd.Values, err = getDMLRouting(upd.Where, eupd.Table)
 	if err != nil {
-		return nil, err
+		eupd.Opcode = engine.UpdateScatter
+	} else {
+		eupd.Opcode = engine.UpdateEqual
+	}
+
+	if eupd.Opcode == engine.UpdateScatter {
+		if len(eupd.Table.Owned) != 0 {
+			return eupd, errors.New("unsupported: multi shard update on a table with owned lookup vindexes")
+		}
+		if upd.Limit != nil {
+			return eupd, errors.New("unsupported: multi shard update with limit")
+		}
 	}
-	eupd.Opcode = engine.UpdateEqual
 
 	if eupd.ChangedVindexValues, err = buildChangedVindexesValues(eupd, upd, eupd.Table.ColumnVindexes); err != nil {
 		return nil, err

test/vtgatev3_test.py

@@ -894,7 +894,7 @@ def test_user2(self):
         {'id': 3})
     vtgate_conn.commit()
 
-    # Test multi shard delete
+    # Test scatter delete
     vtgate_conn.begin()
     result = self.execute_on_master(
         vtgate_conn,
@@ -908,6 +908,20 @@ def test_user2(self):
     self.assertEqual(result, ([], 2L, 0L, []))
     vtgate_conn.commit()
 
+    # Test scatter update
+    vtgate_conn.begin()
+    result = self.execute_on_master(
+        vtgate_conn,
+        'insert into vt_user (id, name) values (:id0, :name0),(:id1, :name1)',
+        {'id0': 22, 'name0': 'name2', 'id1': 33, 'name1': 'name2'})
+    self.assertEqual(result, ([], 2L, 0L, []))
+    result = self.execute_on_master(
+        vtgate_conn,
+        'update vt_user set name=:name where id > :id',
+        {'id': 20, 'name': 'jose'})
+    self.assertEqual(result, ([], 2L, 0L, []))
+    vtgate_conn.commit()
+
   def test_user_truncate(self):
     vtgate_conn = get_connection()
     vtgate_conn.begin()