go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/bigquery/query.go (about)

     1  // Copyright 2021 The Fuchsia Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	"encoding/json"
    10  	"errors"
    11  	"fmt"
    12  	"io"
    13  	"os"
    14  	"strings"
    15  
    16  	"cloud.google.com/go/bigquery"
    17  	"github.com/maruel/subcommands"
    18  	"go.chromium.org/luci/auth"
    19  	"google.golang.org/api/iterator"
    20  	"google.golang.org/api/option"
    21  )
    22  
    23  const queryLongDesc = `Run a SQL query and return the results as JSON.
    24  
    25  Value types will be preserved to the extent that JSON. Booleans and strings
    26  produced by the query will become JSON booleans and strings, respectively,
    27  while ints and floats from the query will both become JSON floats.
    28  `
    29  
    30  func cmdQuery(authOpts auth.Options) *subcommands.Command {
    31  	return &subcommands.Command{
    32  		UsageLine: "query -input <sql-query-file> -json-output <output-path>",
    33  		ShortDesc: "Run a SQL query.",
    34  		LongDesc:  queryLongDesc,
    35  		CommandRun: func() subcommands.CommandRun {
    36  			c := &queryCmd{}
    37  			c.Init(authOpts)
    38  			return c
    39  		},
    40  	}
    41  }
    42  
    43  type queryRow map[string]bigquery.Value
    44  
    45  type queryCmd struct {
    46  	commonFlags
    47  
    48  	inputPath      string
    49  	jsonOutputPath string
    50  }
    51  
    52  func (c *queryCmd) Init(defaultAuthOpts auth.Options) {
    53  	c.commonFlags.Init(defaultAuthOpts)
    54  	c.Flags.StringVar(&c.inputPath, "input", "", "Path to an input file containing an SQL query to run.")
    55  	c.Flags.StringVar(&c.jsonOutputPath, "json-output", "", "Path to output file. Prints to stdout if unspecified.")
    56  }
    57  
    58  func (c *queryCmd) parseArgs() error {
    59  	if err := c.commonFlags.Parse(); err != nil {
    60  		return err
    61  	}
    62  	if c.inputPath == "" {
    63  		return errors.New("-input is required")
    64  	}
    65  	return nil
    66  }
    67  
    68  func (c *queryCmd) Run(a subcommands.Application, _ []string, _ subcommands.Env) int {
    69  	if err := c.parseArgs(); err != nil {
    70  		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
    71  		return 1
    72  	}
    73  
    74  	if err := c.main(); err != nil {
    75  		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
    76  		return 1
    77  	}
    78  	return 0
    79  }
    80  
    81  func (c *queryCmd) main() error {
    82  	ctx := context.Background()
    83  
    84  	authenticator := auth.NewAuthenticator(ctx, auth.OptionalLogin, c.parsedAuthOpts)
    85  	tokenSource, err := authenticator.TokenSource()
    86  	if err != nil {
    87  		if err == auth.ErrLoginRequired {
    88  			fmt.Fprintf(os.Stderr, "You need to login first by running:\n")
    89  			fmt.Fprintf(os.Stderr, "  luci-auth login -scopes %q\n", strings.Join(c.parsedAuthOpts.Scopes, " "))
    90  		}
    91  		return err
    92  	}
    93  
    94  	client, err := bigquery.NewClient(ctx, c.project, option.WithTokenSource(tokenSource))
    95  	if err != nil {
    96  		return err
    97  	}
    98  
    99  	queryBytes, err := os.ReadFile(c.inputPath)
   100  	if err != nil {
   101  		return err
   102  	}
   103  
   104  	rows, err := runQuery(ctx, client, string(queryBytes))
   105  	if err != nil {
   106  		return err
   107  	}
   108  
   109  	return writeOutput(rows, c.jsonOutputPath, os.Stdout)
   110  }
   111  
   112  func runQuery(ctx context.Context, client *bigquery.Client, query string) ([]map[string]bigquery.Value, error) {
   113  	q := client.Query(query)
   114  	iter, err := q.Read(ctx)
   115  	if err != nil {
   116  		return nil, err
   117  	}
   118  
   119  	var rows []map[string]bigquery.Value
   120  	for {
   121  		row := make(map[string]bigquery.Value)
   122  		err := iter.Next(&row)
   123  		if err == iterator.Done {
   124  			break
   125  		} else if err != nil {
   126  			return nil, err
   127  		}
   128  		rows = append(rows, row)
   129  	}
   130  
   131  	return rows, nil
   132  }
   133  
   134  // writeOutput writes the rows returned by a query to the file specified by
   135  // `jsonOutputPath`, or to stdout if `jsonOutputPath` is unset. Takes stdout as
   136  // a parameter to make testing easier.
   137  func writeOutput(rows []map[string]bigquery.Value, jsonOutputPath string, stdout io.Writer) error {
   138  	var output io.Writer
   139  	var indent string
   140  	if jsonOutputPath == "" {
   141  		output = stdout
   142  		// Indent output if printing to stdout for a more user-friendly
   143  		// experience.
   144  		indent = "  "
   145  	} else {
   146  		f, err := os.Create(jsonOutputPath)
   147  		if err != nil {
   148  			return err
   149  		}
   150  		defer f.Close()
   151  		output = f
   152  	}
   153  
   154  	enc := json.NewEncoder(output)
   155  	enc.SetIndent("", indent)
   156  	return enc.Encode(rows)
   157  }