go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/resultdb/cli/cmd_rpc.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package cli
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"io"
    21  	"os"
    22  
    23  	"github.com/maruel/subcommands"
    24  	"google.golang.org/grpc/metadata"
    25  
    26  	"go.chromium.org/luci/auth"
    27  	"go.chromium.org/luci/common/cli"
    28  	"go.chromium.org/luci/common/data/text"
    29  	"go.chromium.org/luci/common/errors"
    30  	"go.chromium.org/luci/grpc/prpc"
    31  )
    32  
    33  const rpcUsage = `rpc [flags] SERVICE METHOD`
    34  
    35  func cmdRPC(p Params) *subcommands.Command {
    36  	return &subcommands.Command{
    37  		UsageLine: rpcUsage,
    38  		ShortDesc: "Make a ResultDB RPC",
    39  		LongDesc: text.Doc(`
    40  			Make a ResultDB RPC.
    41  
    42  			SERVICE must be the full name of a service, e.g. "luci.resultdb.v1.ResultDB"".
    43  			METHOD is the name of the method, e.g. "GetInvocation"
    44  
    45  			The request message is read from stdin, in JSON format.
    46  			The response is printed to stdout, also in JSON format.
    47  		`),
    48  		Advanced: true,
    49  		CommandRun: func() subcommands.CommandRun {
    50  			r := &rpcRun{}
    51  			r.RegisterGlobalFlags(p)
    52  			r.Flags.BoolVar(&r.includeUpdateToken, "include-update-token", false, "send the request with the current invocation's update token in LUCI_CONTEXT")
    53  			return r
    54  		},
    55  	}
    56  }
    57  
    58  type rpcRun struct {
    59  	baseCommandRun
    60  	service            string
    61  	method             string
    62  	includeUpdateToken bool
    63  }
    64  
    65  func (r *rpcRun) parseArgs(args []string) error {
    66  	if len(args) != 2 {
    67  		return errors.Reason("usage: %s", rpcUsage).Err()
    68  	}
    69  
    70  	r.service = args[0]
    71  	r.method = args[1]
    72  
    73  	return nil
    74  }
    75  
    76  func (r *rpcRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
    77  	ctx := cli.GetContext(a, r, env)
    78  
    79  	if err := r.parseArgs(args); err != nil {
    80  		return r.done(err)
    81  	}
    82  
    83  	if err := r.initClients(ctx, auth.SilentLogin); err != nil {
    84  		return r.done(err)
    85  	}
    86  
    87  	if err := r.rpc(ctx); err != nil {
    88  		return r.done(err)
    89  	}
    90  
    91  	return 0
    92  }
    93  
    94  func (r *rpcRun) rpc(ctx context.Context) error {
    95  	// Prepare arguments.
    96  	in, err := io.ReadAll(os.Stdin)
    97  	if err != nil {
    98  		return err
    99  	}
   100  
   101  	if r.includeUpdateToken {
   102  		if r.resultdbCtx == nil {
   103  			return errors.Reason("update token not found, resultdb section of LUCI_CONTEXT missing").Err()
   104  		}
   105  		if r.resultdbCtx.CurrentInvocation.UpdateToken == "" {
   106  			return errors.Reason("update token not found, missing from LUCI_CONTEXT").Err()
   107  		}
   108  		ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("update-token", r.resultdbCtx.CurrentInvocation.UpdateToken))
   109  	}
   110  
   111  	// Send the request.
   112  	res, err := r.prpcClient.CallWithFormats(ctx, r.service, r.method, in, prpc.FormatJSONPB, prpc.FormatJSONPB)
   113  	if err != nil {
   114  		return err
   115  	}
   116  
   117  	// Read response.
   118  	if _, err := os.Stdout.Write(res); err != nil {
   119  		return fmt.Errorf("failed to write response: %s", err)
   120  	}
   121  
   122  	return nil
   123  }