cmd/bigquery/query.go

// Copyright 2021 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"cloud.google.com/go/bigquery"
	"github.com/maruel/subcommands"
	"go.chromium.org/luci/auth"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

const queryLongDesc = `Run a SQL query and return the results as JSON.

Value types will be preserved to the extent that JSON. Booleans and strings
produced by the query will become JSON booleans and strings, respectively,
while ints and floats from the query will both become JSON floats.
`

func cmdQuery(authOpts auth.Options) *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "query -input <sql-query-file> -json-output <output-path>",
		ShortDesc: "Run a SQL query.",
		LongDesc:  queryLongDesc,
		CommandRun: func() subcommands.CommandRun {
			c := &queryCmd{}
			c.Init(authOpts)
			return c
		},
	}
}

type queryRow map[string]bigquery.Value

type queryCmd struct {
	commonFlags

	inputPath      string
	jsonOutputPath string
}

func (c *queryCmd) Init(defaultAuthOpts auth.Options) {
	c.commonFlags.Init(defaultAuthOpts)
	c.Flags.StringVar(&c.inputPath, "input", "", "Path to an input file containing an SQL query to run.")
	c.Flags.StringVar(&c.jsonOutputPath, "json-output", "", "Path to output file. Prints to stdout if unspecified.")
}

func (c *queryCmd) parseArgs() error {
	if err := c.commonFlags.Parse(); err != nil {
		return err
	}
	if c.inputPath == "" {
		return errors.New("-input is required")
	}
	return nil
}

func (c *queryCmd) Run(a subcommands.Application, _ []string, _ subcommands.Env) int {
	if err := c.parseArgs(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	if err := c.main(); err != nil {
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	return 0
}

func (c *queryCmd) main() error {
	ctx := context.Background()

	authenticator := auth.NewAuthenticator(ctx, auth.OptionalLogin, c.parsedAuthOpts)
	tokenSource, err := authenticator.TokenSource()
	if err != nil {
		if err == auth.ErrLoginRequired {
			fmt.Fprintf(os.Stderr, "You need to login first by running:\n")
			fmt.Fprintf(os.Stderr, "  luci-auth login -scopes %q\n", strings.Join(c.parsedAuthOpts.Scopes, " "))
		}
		return err
	}

	client, err := bigquery.NewClient(ctx, c.project, option.WithTokenSource(tokenSource))
	if err != nil {
		return err
	}

	queryBytes, err := os.ReadFile(c.inputPath)
	if err != nil {
		return err
	}

	rows, err := runQuery(ctx, client, string(queryBytes))
	if err != nil {
		return err
	}

	return writeOutput(rows, c.jsonOutputPath, os.Stdout)
}

func runQuery(ctx context.Context, client *bigquery.Client, query string) ([]map[string]bigquery.Value, error) {
	q := client.Query(query)
	iter, err := q.Read(ctx)
	if err != nil {
		return nil, err
	}

	var rows []map[string]bigquery.Value
	for {
		row := make(map[string]bigquery.Value)
		err := iter.Next(&row)
		if err == iterator.Done {
			break
		} else if err != nil {
			return nil, err
		}
		rows = append(rows, row)
	}

	return rows, nil
}

// writeOutput writes the rows returned by a query to the file specified by
// `jsonOutputPath`, or to stdout if `jsonOutputPath` is unset. Takes stdout as
// a parameter to make testing easier.
func writeOutput(rows []map[string]bigquery.Value, jsonOutputPath string, stdout io.Writer) error {
	var output io.Writer
	var indent string
	if jsonOutputPath == "" {
		output = stdout
		// Indent output if printing to stdout for a more user-friendly
		// experience.
		indent = "  "
	} else {
		f, err := os.Create(jsonOutputPath)
		if err != nil {
			return err
		}
		defer f.Close()
		output = f
	}

	enc := json.NewEncoder(output)
	enc.SetIndent("", indent)
	return enc.Encode(rows)
}