v3: protect from partial DML execution

V3 sometimes breaks up a single DML into multiple statements.
Sometimes, this may lead to failures in the middle. The code
change detects such failures and rolls back the transaction
to prevent commits that are accidentally partial.

Author: sougou

go/vt/vtgate/engine/primitive.go

@@ -22,8 +22,8 @@ const ListVarName = "__vals"
 // VCursor defines the interface the engine will use
 // to execute routes.
 type VCursor interface {
-	Execute(query string, bindvars map[string]interface{}) (*sqltypes.Result, error)
-	ExecuteMultiShard(keyspace string, shardQueries map[string]querytypes.BoundQuery) (*sqltypes.Result, error)
+	Execute(query string, bindvars map[string]interface{}, isDML bool) (*sqltypes.Result, error)
+	ExecuteMultiShard(keyspace string, shardQueries map[string]querytypes.BoundQuery, isDML bool) (*sqltypes.Result, error)
 	ExecuteStandalone(query string, bindvars map[string]interface{}, keyspace, shard string) (*sqltypes.Result, error)
 	StreamExecuteMulti(query string, keyspace string, shardVars map[string]map[string]interface{}, callback func(reply *sqltypes.Result) error) error
 	GetKeyspaceShards(vkeyspace *vindexes.Keyspace) (string, []*topodatapb.ShardReference, error)

go/vt/vtgate/engine/route.go

@@ -257,8 +257,12 @@ func (route *Route) Execute(vcursor VCursor, bindVars, joinVars map[string]inter
 
 	var err error
 	var params *scatterParams
+	isDML := false
 	switch route.Opcode {
-	case SelectUnsharded, UpdateUnsharded, DeleteUnsharded, SelectScatter:
+	case SelectUnsharded, SelectScatter:
+		params, err = route.paramsAllShards(vcursor, bindVars)
+	case UpdateUnsharded, DeleteUnsharded:
+		isDML = true
 		params, err = route.paramsAllShards(vcursor, bindVars)
 	case SelectEqual, SelectEqualUnique:
 		params, err = route.paramsSelectEqual(vcursor, bindVars)
@@ -273,7 +277,7 @@ func (route *Route) Execute(vcursor VCursor, bindVars, joinVars map[string]inter
 	}
 
 	shardQueries := route.getShardQueries(route.Query, params)
-	return vcursor.ExecuteMultiShard(params.ks, shardQueries)
+	return vcursor.ExecuteMultiShard(params.ks, shardQueries, isDML)
 }
 
 // StreamExecute performs a streaming exec.
@@ -312,7 +316,7 @@ func (route *Route) GetFields(vcursor VCursor, bindVars, joinVars map[string]int
 		return nil, err
 	}
 
-	return route.execShard(vcursor, route.FieldQuery, bindVars, ks, shard)
+	return route.execShard(vcursor, route.FieldQuery, bindVars, ks, shard, false /* isDML */)
 }
 
 func combineVars(bv1, bv2 map[string]interface{}) map[string]interface{} {
@@ -390,7 +394,7 @@ func (route *Route) execUpdateEqual(vcursor VCursor, bindVars map[string]interfa
 		return &sqltypes.Result{}, nil
 	}
 	rewritten := sqlannotation.AddKeyspaceIDs(route.Query, [][]byte{ksid}, "")
-	return route.execShard(vcursor, rewritten, bindVars, ks, shard)
+	return route.execShard(vcursor, rewritten, bindVars, ks, shard, true /* isDML */)
 }
 
 func (route *Route) execDeleteEqual(vcursor VCursor, bindVars map[string]interface{}) (*sqltypes.Result, error) {
@@ -412,7 +416,7 @@ func (route *Route) execDeleteEqual(vcursor VCursor, bindVars map[string]interfa
 		}
 	}
 	rewritten := sqlannotation.AddKeyspaceIDs(route.Query, [][]byte{ksid}, "")
-	return route.execShard(vcursor, rewritten, bindVars, ks, shard)
+	return route.execShard(vcursor, rewritten, bindVars, ks, shard, true /* isDML */)
 }
 
 func (route *Route) execInsertUnsharded(vcursor VCursor, bindVars map[string]interface{}) (*sqltypes.Result, error) {
@@ -426,7 +430,7 @@ func (route *Route) execInsertUnsharded(vcursor VCursor, bindVars map[string]int
 	}
 
 	shardQueries := route.getShardQueries(route.Query, params)
-	result, err := vcursor.ExecuteMultiShard(params.ks, shardQueries)
+	result, err := vcursor.ExecuteMultiShard(params.ks, shardQueries, true /* isDML */)
 	if err != nil {
 		return nil, fmt.Errorf("execInsertUnsharded: %v", err)
 	}
@@ -451,7 +455,7 @@ func (route *Route) execInsertSharded(vcursor VCursor, bindVars map[string]inter
 		return nil, fmt.Errorf("execInsertSharded: %v", err)
 	}
 
-	result, err := vcursor.ExecuteMultiShard(keyspace, shardQueries)
+	result, err := vcursor.ExecuteMultiShard(keyspace, shardQueries, true /* isDML */)
 
 	if err != nil {
 		return nil, fmt.Errorf("execInsertSharded: %v", err)
@@ -640,7 +644,7 @@ func (route *Route) resolveSingleShard(vcursor VCursor, bindVars map[string]inte
 }
 
 func (route *Route) deleteVindexEntries(vcursor VCursor, bindVars map[string]interface{}, ks, shard string, ksid []byte) error {
-	result, err := route.execShard(vcursor, route.Subquery, bindVars, ks, shard)
+	result, err := route.execShard(vcursor, route.Subquery, bindVars, ks, shard, false /* isDML */)
 	if err != nil {
 		return err
 	}
@@ -807,13 +811,13 @@ func (route *Route) handleNonPrimary(vcursor VCursor, vindexKeys []interface{},
 	return nil
 }
 
-func (route *Route) execShard(vcursor VCursor, query string, bindVars map[string]interface{}, keyspace, shard string) (*sqltypes.Result, error) {
+func (route *Route) execShard(vcursor VCursor, query string, bindVars map[string]interface{}, keyspace, shard string, isDML bool) (*sqltypes.Result, error) {
 	return vcursor.ExecuteMultiShard(keyspace, map[string]querytypes.BoundQuery{
 		shard: {
 			Sql:           query,
 			BindVariables: bindVars,
 		},
-	})
+	}, isDML)
 }
 
 func (route *Route) anyShard(vcursor VCursor, keyspace *vindexes.Keyspace) (string, string, error) {

go/vt/vtgate/executor.go

@@ -145,7 +145,13 @@ func (e *Executor) handleExec(ctx context.Context, session *vtgatepb.Session, sq
 	if err != nil {
 		return nil, err
 	}
-	return plan.Instructions.Execute(vcursor, bindVars, make(map[string]interface{}), true)
+	qr, err := plan.Instructions.Execute(vcursor, bindVars, make(map[string]interface{}), true)
+	// Check if there was partial DML execution. If so, rollback the transaction.
+	if err != nil && session.InTransaction && vcursor.hasPartialDML {
+		_ = e.txConn.Rollback(ctx, NewSafeSession(session))
+		err = vterrors.Errorf(vtrpcpb.Code_ABORTED, "transaction rolled back due to partial DML execution: %v", err)
+	}
+	return qr, err
 }
 
 func (e *Executor) shardExec(ctx context.Context, session *vtgatepb.Session, sql string, bindVars map[string]interface{}, target querypb.Target) (*sqltypes.Result, error) {

go/vt/vtgate/executor_dml_test.go

@@ -5,6 +5,7 @@
 package vtgate
 
 import (
+	"context"
 	"reflect"
 	"strings"
 	"testing"
@@ -15,6 +16,7 @@ import (
 	"github.com/youtube/vitess/go/vt/vttablet/tabletserver/querytypes"
 
 	querypb "github.com/youtube/vitess/go/vt/proto/query"
+	vtgatepb "github.com/youtube/vitess/go/vt/proto/vtgate"
 	vtrpcpb "github.com/youtube/vitess/go/vt/proto/vtrpc"
 )
 
@@ -785,6 +787,36 @@ func TestInsertFail(t *testing.T) {
 	}
 }
 
+func TestInsertPartialFail(t *testing.T) {
+	executor, sbc1, _, sbclookup := createExecutorEnv()
+
+	// If the first DML fails, there should be no rollback.
+	sbclookup.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
+	_, err := executor.Execute(
+		context.Background(),
+		&vtgatepb.Session{InTransaction: true},
+		"insert into user(id, v, name) values (1, 2, 'myname')",
+		nil,
+	)
+	want := "execInsertSharded:"
+	if err == nil || !strings.HasPrefix(err.Error(), want) {
+		t.Errorf("insert first DML fail: %v, must start with %s", err, want)
+	}
+
+	// If the second DML fails, we should rollback.
+	sbc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
+	_, err = executor.Execute(
+		context.Background(),
+		&vtgatepb.Session{InTransaction: true},
+		"insert into user(id, v, name) values (1, 2, 'myname')",
+		nil,
+	)
+	want = "transaction rolled back"
+	if err == nil || !strings.HasPrefix(err.Error(), want) {
+		t.Errorf("insert first DML fail: %v, must start with %s", err, want)
+	}
+}
+
 func TestMultiInsertSharded(t *testing.T) {
 	executor, sbc1, sbc2, sbclookup := createExecutorEnv()
 

go/vt/vtgate/vcursor_impl.go

@@ -27,6 +27,10 @@ type vcursorImpl struct {
 	target           querypb.Target
 	trailingComments string
 	executor         *Executor
+	// hasPartialDML is set to true if any DML was successfully
+	// executed. If there was a subsequent failure, the transaction
+	// must be forced to rollback.
+	hasPartialDML bool
 }
 
 // newVcursorImpl creates a vcursorImpl. Before creating this object, you have to separate out any trailingComments that came with
@@ -54,13 +58,21 @@ func (vc *vcursorImpl) Find(keyspace, tablename sqlparser.TableIdent) (table *vi
 }
 
 // Execute performs a V3 level execution of the query. It does not take any routing directives.
-func (vc *vcursorImpl) Execute(query string, BindVars map[string]interface{}) (*sqltypes.Result, error) {
-	return vc.executor.Execute(vc.ctx, vc.session, query+vc.trailingComments, BindVars)
+func (vc *vcursorImpl) Execute(query string, BindVars map[string]interface{}, isDML bool) (*sqltypes.Result, error) {
+	qr, err := vc.executor.Execute(vc.ctx, vc.session, query+vc.trailingComments, BindVars)
+	if err == nil {
+		vc.hasPartialDML = true
+	}
+	return qr, err
 }
 
 // ExecuteMultiShard executes different queries on different shards and returns the combined result.
-func (vc *vcursorImpl) ExecuteMultiShard(keyspace string, shardQueries map[string]querytypes.BoundQuery) (*sqltypes.Result, error) {
-	return vc.executor.scatterConn.ExecuteMultiShard(vc.ctx, keyspace, commentedShardQueries(shardQueries, vc.trailingComments), vc.target.TabletType, NewSafeSession(vc.session), false, vc.session.Options)
+func (vc *vcursorImpl) ExecuteMultiShard(keyspace string, shardQueries map[string]querytypes.BoundQuery, isDML bool) (*sqltypes.Result, error) {
+	qr, err := vc.executor.scatterConn.ExecuteMultiShard(vc.ctx, keyspace, commentedShardQueries(shardQueries, vc.trailingComments), vc.target.TabletType, NewSafeSession(vc.session), false, vc.session.Options)
+	if err == nil {
+		vc.hasPartialDML = true
+	}
+	return qr, err
 }
 
 // ExecuteStandalone executes the specified query on keyspace:shard, but outside of the current transaction, as an independent statement.
@@ -71,7 +83,11 @@ func (vc *vcursorImpl) ExecuteStandalone(query string, BindVars map[string]inter
 			BindVariables: BindVars,
 		},
 	}
-	return vc.executor.scatterConn.ExecuteMultiShard(vc.ctx, keyspace, bq, vc.target.TabletType, NewSafeSession(nil), false, vc.session.Options)
+	qr, err := vc.executor.scatterConn.ExecuteMultiShard(vc.ctx, keyspace, bq, vc.target.TabletType, NewSafeSession(nil), false, vc.session.Options)
+	if err == nil {
+		vc.hasPartialDML = true
+	}
+	return qr, err
 }
 
 // StreamExeculteMulti is the streaming version of ExecuteMultiShard.

go/vt/vtgate/vindexes/lookup_hash_test.go

@@ -23,7 +23,7 @@ type vcursor struct {
 	bq       *querytypes.BoundQuery
 }
 
-func (vc *vcursor) Execute(query string, bindvars map[string]interface{}) (*sqltypes.Result, error) {
+func (vc *vcursor) Execute(query string, bindvars map[string]interface{}, isDML bool) (*sqltypes.Result, error) {
 	vc.bq = &querytypes.BoundQuery{
 		Sql:           query,
 		BindVariables: bindvars,

go/vt/vtgate/vindexes/lookup_internal.go

@@ -37,7 +37,7 @@ func (lkp *lookup) MapUniqueLookup(vcursor VCursor, ids []interface{}) ([][]byte
 	for _, id := range ids {
 		result, err := vcursor.Execute(lkp.sel, map[string]interface{}{
 			lkp.From: id,
-		})
+		}, false /* isDML */)
 		if err != nil {
 			return nil, fmt.Errorf("lookup.Map: %v", err)
 		}
@@ -67,7 +67,7 @@ func (lkp *lookup) MapNonUniqueLookup(vcursor VCursor, ids []interface{}) ([][][
 	for _, id := range ids {
 		result, err := vcursor.Execute(lkp.sel, map[string]interface{}{
 			lkp.From: id,
-		})
+		}, false /* isDML */)
 		if err != nil {
 			return nil, fmt.Errorf("lookup.Map: %v", err)
 		}
@@ -124,7 +124,7 @@ func (lkp *lookup) Verify(vcursor VCursor, ids []interface{}, ksids [][]byte) (b
 		bindVars[toStr] = val[rowNum]
 	}
 	lkp.ver = fmt.Sprintf("select %s from %s where %s", lkp.From, lkp.Table, strings.Trim(colBuff.String(), "or")+")")
-	result, err := vcursor.Execute(lkp.ver, bindVars)
+	result, err := vcursor.Execute(lkp.ver, bindVars, false /* isDML */)
 	if err != nil {
 		return false, fmt.Errorf("lookup.Verify: %v", err)
 	}
@@ -168,7 +168,7 @@ func (lkp *lookup) Create(vcursor VCursor, ids []interface{}, ksids [][]byte) er
 		bindVars[toStr] = val[rowNum]
 	}
 	lkp.ins = strings.Trim(insBuffer.String(), ",")
-	if _, err := vcursor.Execute(lkp.ins, bindVars); err != nil {
+	if _, err := vcursor.Execute(lkp.ins, bindVars, true /* isDML */); err != nil {
 		return fmt.Errorf("lookup.Create: %v", err)
 	}
 	return nil
@@ -191,7 +191,7 @@ func (lkp *lookup) Delete(vcursor VCursor, ids []interface{}, ksid []byte) error
 	}
 	for _, id := range ids {
 		bindvars[lkp.From] = id
-		if _, err := vcursor.Execute(lkp.del, bindvars); err != nil {
+		if _, err := vcursor.Execute(lkp.del, bindvars, true /* isDML */); err != nil {
 			return fmt.Errorf("lookup.Delete: %v", err)
 		}
 	}

go/vt/vtgate/vindexes/vindex.go

@@ -16,7 +16,7 @@ import (
 // in the current context and session of a VTGate request. Vindexes
 // can use this interface to execute lookup queries.
 type VCursor interface {
-	Execute(query string, bindvars map[string]interface{}) (*sqltypes.Result, error)
+	Execute(query string, bindvars map[string]interface{}, isDML bool) (*sqltypes.Result, error)
 }
 
 // Vindex defines the interface required to register a vindex.