Applyschema uses ExecuteMultiFetchAsDba

Signed-off-by: Shlomi Noach <[email protected]>

go/vt/schemamanager/local_controller.go

@@ -208,8 +208,10 @@ func (controller *LocalController) writeToLogDir(ctx context.Context, result *Ex
 	rowsReturned := uint64(0)
 	rowsAffected := uint64(0)
 	for _, queryResult := range result.SuccessShards {
-		rowsReturned += uint64(len(queryResult.Result.Rows))
-		rowsAffected += queryResult.Result.RowsAffected
+		for _, result := range queryResult.Results {
+			rowsReturned += uint64(len(result.Rows))
+			rowsAffected += result.RowsAffected
+		}
 	}
 	logFile.WriteString(fmt.Sprintf("-- Rows returned: %d\n", rowsReturned))
 	logFile.WriteString(fmt.Sprintf("-- Rows affected: %d\n", rowsAffected))

go/vt/schemamanager/local_controller_test.go

@@ -138,8 +138,8 @@ func TestLocalControllerSchemaChange(t *testing.T) {
 	result := &ExecuteResult{
 		Sqls: []string{"create table test_table (id int)"},
 		SuccessShards: []ShardResult{{
-			Shard:  "0",
-			Result: &querypb.QueryResult{},
+			Shard:   "0",
+			Results: []*querypb.QueryResult{{}},
 		}},
 	}
 	logPath := path.Join(controller.logDir, controller.sqlFilename)

go/vt/schemamanager/schemamanager.go

@@ -85,8 +85,8 @@ type ShardWithError struct {
 
 // ShardResult contains sql execute information on a particular shard
 type ShardResult struct {
-	Shard  string
-	Result *querypb.QueryResult
+	Shard   string
+	Results []*querypb.QueryResult
 	// Position is a replication position that is guaranteed to be after the
 	// schema change was applied. It can be used to wait for replicas to receive
 	// the schema change via replication.

go/vt/schemamanager/tablet_executor.go

@@ -572,13 +572,15 @@ func (exec *TabletExecutor) executeOneTablet(
 	errChan chan ShardWithError,
 	successChan chan ShardResult) {
 
-	var result *querypb.QueryResult
+	var results []*querypb.QueryResult
 	var err error
 	if viaQueryService {
-		result, err = exec.tmc.ExecuteQuery(ctx, tablet, &tabletmanagerdatapb.ExecuteQueryRequest{
+		result, reserr := exec.tmc.ExecuteQuery(ctx, tablet, &tabletmanagerdatapb.ExecuteQueryRequest{
 			Query:   []byte(sql),
 			MaxRows: 10,
 		})
+		results = []*querypb.QueryResult{result}
+		err = reserr
 	} else {
 		if exec.ddlStrategySetting != nil && exec.ddlStrategySetting.IsAllowZeroInDateFlag() {
 			// --allow-zero-in-date Applies to DDLs
@@ -588,14 +590,14 @@ func (exec *TabletExecutor) executeOneTablet(
 				return
 			}
 		}
-		request := &tabletmanagerdatapb.ExecuteFetchAsDbaRequest{
-			Query:   []byte(sql),
+		request := &tabletmanagerdatapb.ExecuteMultiFetchAsDbaRequest{
+			Sql:     []byte(sql),
 			MaxRows: 10,
 		}
 		if exec.ddlStrategySetting != nil && exec.ddlStrategySetting.IsAllowForeignKeysFlag() {
 			request.DisableForeignKeyChecks = true
 		}
-		result, err = exec.tmc.ExecuteFetchAsDba(ctx, tablet, false, request)
+		results, err = exec.tmc.ExecuteMultiFetchAsDba(ctx, tablet, false, request)
 
 	}
 	if err != nil {
@@ -614,7 +616,7 @@ func (exec *TabletExecutor) executeOneTablet(
 	}
 	successChan <- ShardResult{
 		Shard:    tablet.Shard,
-		Result:   result,
+		Results:  results,
 		Position: pos,
 	}
 }

go/vt/vtctl/grpcvtctldserver/server.go

@@ -303,7 +303,9 @@ func (s *VtctldServer) ApplySchema(ctx context.Context, req *vtctldatapb.ApplySc
 	}
 
 	for _, shard := range execResult.SuccessShards {
-		resp.RowsAffectedByShard[shard.Shard] = shard.Result.RowsAffected
+		for _, result := range shard.Results {
+			resp.RowsAffectedByShard[shard.Shard] += result.RowsAffected
+		}
 	}
 
 	return resp, err