vtgate API: test AsTransaction

go/vt/vtgate/resolver_test.go

@@ -512,6 +512,51 @@ func TestResolverDmlOnMultipleKeyspaceIds(t *testing.T) {
 	}
 }
 
+func TestResolverExecBatchAsTransaction(t *testing.T) {
+	s := createSandbox(KsTestUnshardedServedFrom)
+	sbc := &sandboxConn{mustFailRetry: 20}
+	s.MapTestConn("0", sbc)
+
+	res := NewResolver(new(sandboxTopo), "", "aa", 1*time.Millisecond, 2, 2*time.Millisecond, 1*time.Millisecond, 24*time.Hour)
+
+	callcount := 0
+	buildBatchRequest := func() (*scatterBatchRequest, error) {
+		callcount++
+		queries := []proto.BoundShardQuery{{
+			Sql:           "query",
+			BindVariables: nil,
+			Keyspace:      KsTestUnshardedServedFrom,
+			Shards:        []string{"0"},
+		}}
+		return boundShardQueriesToScatterBatchRequest(queries), nil
+	}
+
+	_, err := res.ExecuteBatch(context.Background(), topo.TYPE_MASTER, false, nil, buildBatchRequest)
+	if err == nil {
+		t.Errorf("want got, got none")
+	}
+	// Ensure scatter tried a re-resolve
+	if callcount != 2 {
+		t.Errorf("want 2, got %v", callcount)
+	}
+	if sbc.AsTransactionCount != 0 {
+		t.Errorf("want 0, got %v", sbc.AsTransactionCount)
+	}
+
+	callcount = 0
+	_, err = res.ExecuteBatch(context.Background(), topo.TYPE_MASTER, true, nil, buildBatchRequest)
+	if err == nil {
+		t.Errorf("want got, got none")
+	}
+	// Ensure scatter did not re-resolve
+	if callcount != 1 {
+		t.Errorf("want 1, got %v", callcount)
+	}
+	if sbc.AsTransactionCount != 1 {
+		t.Errorf("want 1, got %v", sbc.AsTransactionCount)
+	}
+}
+
 func TestIsConnError(t *testing.T) {
 	var connErrorTests = []struct {
 		in      error

go/vt/vtgate/sandbox_test.go

@@ -327,11 +327,12 @@ type sandboxConn struct {
 
 	// These Count vars report how often the corresponding
 	// functions were called.
-	ExecCount     sync2.AtomicInt64
-	BeginCount    sync2.AtomicInt64
-	CommitCount   sync2.AtomicInt64
-	RollbackCount sync2.AtomicInt64
-	CloseCount    sync2.AtomicInt64
+	ExecCount          sync2.AtomicInt64
+	BeginCount         sync2.AtomicInt64
+	CommitCount        sync2.AtomicInt64
+	RollbackCount      sync2.AtomicInt64
+	CloseCount         sync2.AtomicInt64
+	AsTransactionCount sync2.AtomicInt64
 
 	// Queries stores the requests received.
 	Queries []tproto.BoundQuery
@@ -405,6 +406,9 @@ func (sbc *sandboxConn) Execute2(ctx context.Context, query string, bindVars map
 
 func (sbc *sandboxConn) ExecuteBatch(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (*tproto.QueryResultList, error) {
 	sbc.ExecCount.Add(1)
+	if asTransaction {
+		sbc.AsTransactionCount.Add(1)
+	}
 	if sbc.mustDelay != 0 {
 		time.Sleep(sbc.mustDelay)
 	}

go/vt/vtgate/shard_conn.go

@@ -125,7 +125,6 @@ func (sdc *ShardConn) Execute(ctx context.Context, query string, bindVars map[st
 
 // ExecuteBatch executes a group of queries. The retry rules are the same as Execute.
 func (sdc *ShardConn) ExecuteBatch(ctx context.Context, queries []tproto.BoundQuery, asTransaction bool, transactionID int64) (qrs *tproto.QueryResultList, err error) {
-	// FIXME(sougou): don't retry if asTransaction is true.
 	err = sdc.withRetry(ctx, func(conn tabletconn.TabletConn) error {
 		var innerErr error
 		qrs, innerErr = conn.ExecuteBatch(ctx, queries, asTransaction, transactionID)
@@ -223,6 +222,10 @@ func (sdc *ShardConn) withRetry(ctx context.Context, action func(conn tabletconn
 	for i := 0; i < sdc.retryCount+1; i++ {
 		conn, endPoint, isTimeout, err = sdc.getConn(ctx)
 		if err != nil {
+			// CAUTION: if we decide to retry on timeout, remember to only
+			// retry read queries. We now allow DMLs in autocommit mode,
+			// and we also allow AsTransaction flag in ExecuteBatch. In those
+			// cases, we should not retry on timeout.
 			if isTimeout || i == sdc.retryCount {
 				break
 			}