// RunCommandAndWait executes a single command on a given vtworker and blocks until the command did return or timed out.
// Output from vtworker is streamed as logutil.Event messages which
// have to be consumed by the caller who has to specify a "recv" function.
func RunCommandAndWait(ctx context.Context, server string, args []string, recv func(*logutilpb.Event)) error {
if recv == nil {
panic("no function closure for Event stream specified")
}
// create the client
// TODO(mberlin): vtctlclient exposes dialTimeout as flag. If there are no use cases, remove it there as well to be consistent?
client, err := New(server, 30*time.Second /* dialTimeout */)
if err != nil {
return vterrors.WithPrefix("cannot dial to server "+server+": ", err)
}
defer client.Close()
// run the command
stream, err := client.ExecuteVtworkerCommand(ctx, args)
if err != nil {
return vterrors.WithPrefix("cannot execute remote command: ", err)
}
for {
e, err := stream.Recv()
switch err {
case nil:
recv(e)
case io.EOF:
return nil
default:
return vterrors.WithPrefix("stream error: ", err)
}
}
}
// RunCommand executes the vtworker command specified by "args". Use WaitForCommand() to block on the returned done channel.
// If wr is nil, the default wrangler will be used.
// If you pass a wr wrangler, note that a MemoryLogger will be added to its current logger.
// The returned worker and done channel may be nil if no worker was started e.g. in case of a "Reset".
func (wi *Instance) RunCommand(args []string, wr *wrangler.Wrangler, runFromCli bool) (Worker, chan struct{}, error) {
if len(args) >= 1 {
switch args[0] {
case "Reset":
err := wi.Reset()
return nil, nil, err
case "Cancel":
wi.Cancel()
return nil, nil, nil
}
}
if wr == nil {
wr = wi.wr
}
wrk, err := commandWorker(wi, wr, args, wi.cell, runFromCli)
if err != nil {
return nil, nil, err
}
done, err := wi.setAndStartWorker(wrk, wr)
if err != nil {
return nil, nil, vterrors.WithPrefix("cannot set worker: ", err)
}
return wrk, done, nil
}
func runVtworkerCommand(client vtworkerclient.Client, args []string) error {
stream, err := client.ExecuteVtworkerCommand(context.Background(), args)
if err != nil {
return fmt.Errorf("cannot execute remote command: %v", err)
}
for {
_, err := stream.Recv()
switch err {
case nil:
// Consume next response.
case io.EOF:
return nil
default:
return vterrors.WithPrefix("unexpected error when reading the stream: ", err)
}
}
}
