| // Copyright 2018 The Chromium Authors |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| // Command bqupload inserts rows in a BigQuery table. |
| // |
| // It is a lightweight alternative to 'bq insert' command from gcloud SDK. |
| // |
| // Inserts the records formatted as newline delimited JSON from file into |
| // the specified table. If file is not specified, reads from stdin. If there |
| // were any insert errors it prints the errors to stderr. |
| // |
| // Usage: |
| // |
| // bqupload <project>.<dataset>.<table> [<file>] |
| package main |
| |
| import ( |
| "bufio" |
| "bytes" |
| "context" |
| "crypto/rand" |
| "encoding/base64" |
| "encoding/json" |
| "flag" |
| "fmt" |
| "io" |
| "os" |
| "strings" |
| "sync" |
| |
| "cloud.google.com/go/bigquery" |
| "golang.org/x/oauth2" |
| "golang.org/x/sync/errgroup" |
| "google.golang.org/api/option" |
| |
| "go.chromium.org/luci/auth" |
| "go.chromium.org/luci/auth/client/authcli" |
| "go.chromium.org/luci/common/errors" |
| "go.chromium.org/luci/common/flag/stringmapflag" |
| "go.chromium.org/luci/common/logging" |
| "go.chromium.org/luci/common/logging/gologger" |
| "go.chromium.org/luci/hardcoded/chromeinfra" |
| ) |
| |
| const ( |
| userAgent = "bqupload v1.6" |
| // The bigquery API imposes a hard limit of 50,000 rows. We use a much lower |
| // default limit to also make it less likely that the total payload size |
| // exceeds the maximum, and to limit the blast radius when a batch fails to |
| // upload. |
| defaultBatchSize = 500 |
| // crbug/1491321 - http2 concurrent request limits |
| // While using http2 in Golang, try to keep max concurrent requests to 100. |
| maxConcurrentInserts = 100 |
| ) |
| |
| func usage() { |
| fmt.Fprintf(os.Stderr, |
| `%s |
| |
| Usage: bqupload <project>.<dataset>.<table> [<file>] |
| |
| Inserts the records formatted as newline delimited JSON from file into |
| the specified table. If file is not specified, reads from stdin. If there |
| were any insert errors it prints the errors to stderr. |
| |
| Optional flags: |
| `, userAgent) |
| flag.PrintDefaults() |
| } |
| |
| func main() { |
| flag.CommandLine.Usage = usage |
| if err := run(gologger.StdConfig.Use(context.Background())); err != nil { |
| fmt.Fprintf(os.Stderr, "bqupload: %s\n", err) |
| os.Exit(1) |
| } |
| } |
| |
| type uploadOpts struct { |
| project string |
| dataset string |
| table string |
| columns stringmapflag.Value |
| |
| input io.Reader |
| auth oauth2.TokenSource |
| insertIDBase string |
| |
| ignoreUnknownValues bool |
| skipInvalidRows bool |
| jsonList bool |
| batchSize int |
| } |
| |
| func run(ctx context.Context) error { |
| // BQ options. |
| bqOpts := uploadOpts{} |
| flag.BoolVar(&bqOpts.ignoreUnknownValues, "ignore-unknown-values", false, |
| "Ignore any values in a row that are not present in the schema.") |
| flag.BoolVar(&bqOpts.skipInvalidRows, "skip-invalid-rows", false, |
| "Attempt to insert any valid rows, even if invalid rows are present.") |
| flag.IntVar(&bqOpts.batchSize, "batch-size", defaultBatchSize, |
| "Number of rows per insert batch.") |
| flag.BoolVar(&bqOpts.jsonList, "json-list", false, |
| "Instead of looking for newline delimited rows, looks for a JSON list of rows.") |
| flag.Var(&bqOpts.columns, "column", |
| `Parse all the rows as usual, but add or replace these columns in each row. The |
| value is parsed as JSON. Can be specified multiple times to set/replace multiple columns.`) |
| |
| // Auth options. |
| defaults := chromeinfra.DefaultAuthOptions() |
| defaults.Scopes = []string{ |
| "https://www.googleapis.com/auth/bigquery", |
| "https://www.googleapis.com/auth/userinfo.email", |
| } |
| authFlags := authcli.Flags{} |
| authFlags.Register(flag.CommandLine, defaults) |
| |
| flag.Parse() |
| |
| // Parse positional flags. |
| args := flag.Args() |
| if len(args) == 0 || len(args) > 2 { |
| usage() |
| os.Exit(2) |
| } |
| |
| var err error |
| bqOpts.project, bqOpts.dataset, bqOpts.table, err = parseTableRef(args[0]) |
| if err != nil { |
| return err |
| } |
| |
| bqOpts.input = os.Stdin |
| if len(args) > 1 { |
| f, err := os.Open(args[1]) |
| if err != nil { |
| return err |
| } |
| defer f.Close() |
| bqOpts.input = f |
| } |
| |
| // Prepare random prefix to use for insert IDs uploaded by this process. |
| rnd := make([]byte, 12) |
| if _, err = rand.Read(rnd); err != nil { |
| return err |
| } |
| bqOpts.insertIDBase = base64.RawURLEncoding.EncodeToString(rnd) |
| |
| // Get oauth2.TokenSource based on parsed auth flags. |
| authOpts, err := authFlags.Options() |
| if err != nil { |
| return err |
| } |
| authenticator := auth.NewAuthenticator(ctx, auth.SilentLogin, authOpts) |
| bqOpts.auth, 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(defaults.Scopes, " ")) |
| } |
| return err |
| } |
| |
| // Report who we are running as, helps when debugging permissions. Carry on |
| // on errors (there shouldn't be any anyway). |
| email, err := authenticator.GetEmail() |
| if err != nil { |
| logging.Warningf(ctx, "Can't get an email of the active account - %s", err) |
| } else { |
| logging.Infof(ctx, "Running as %s", email) |
| } |
| |
| return upload(ctx, &bqOpts) |
| } |
| |
| func parseTableRef(ref string) (project, dataset, table string, err error) { |
| chunks := strings.Split(ref, ".") |
| if len(chunks) != 3 { |
| err = fmt.Errorf("table reference should have form <project>.<dataset>.<table>, got %q", ref) |
| return |
| } |
| return chunks[0], chunks[1], chunks[2], nil |
| } |
| |
| func upload(ctx context.Context, opts *uploadOpts) error { |
| client, err := bigquery.NewClient(ctx, opts.project, |
| option.WithTokenSource(opts.auth), |
| option.WithUserAgent(userAgent)) |
| if err != nil { |
| return err |
| } |
| defer client.Close() |
| |
| inserter := client.Dataset(opts.dataset).Table(opts.table).Inserter() |
| inserter.IgnoreUnknownValues = opts.ignoreUnknownValues |
| inserter.SkipInvalidRows = opts.skipInvalidRows |
| |
| // Prepare column overrides. |
| overrides := make(map[string]bigquery.Value, len(opts.columns)) |
| for key, value := range opts.columns { |
| var val bigquery.Value |
| if err := json.Unmarshal([]byte(value), &val); err != nil { |
| return errors.Annotate(err, "parsing -column %q value", key).Err() |
| } |
| overrides[key] = val |
| } |
| |
| // Note: we may potentially read rows from 'input' and upload them at the same |
| // time for true streaming uploads in case 'input' is stdin and it's produced |
| // on the fly. This is not trivial though and isn't needed yet, so we read |
| // everything at once. |
| rows, err := readInput(opts.input, opts.insertIDBase, opts.jsonList, overrides) |
| if err != nil { |
| return err |
| } |
| |
| logging.Infof(ctx, |
| "Inserting %d rows into table `%s.%s.%s`", |
| len(rows), opts.project, opts.dataset, opts.table) |
| |
| if err := doInsert(ctx, os.Stderr, opts, inserter, rows); err != nil { |
| return err |
| } |
| |
| logging.Infof(ctx, "Done") |
| return nil |
| } |
| |
| // For testability. |
| type bqInserter interface { |
| Put(ctx context.Context, src interface{}) error |
| } |
| |
| func doInsert(ctx context.Context, stderr io.Writer, opts *uploadOpts, inserter bqInserter, rows []*tableRow) error { |
| var mu sync.Mutex |
| var multiErr bigquery.PutMultiError |
| eg, egCtx := errgroup.WithContext(ctx) |
| eg.SetLimit(maxConcurrentInserts) |
| for i := 0; i < len(rows); i += opts.batchSize { |
| i := i |
| eg.Go(func() error { |
| end := i + opts.batchSize |
| if end > len(rows) { |
| end = len(rows) |
| } |
| batch := rows[i:end] |
| if err := inserter.Put(egCtx, batch); err != nil { |
| if merr, ok := err.(bigquery.PutMultiError); ok { |
| mu.Lock() |
| multiErr = append(multiErr, merr...) |
| mu.Unlock() |
| } else { |
| // Unrecognized errors should be considered fatal. |
| return err |
| } |
| } |
| return nil |
| }) |
| } |
| if err := eg.Wait(); err != nil { |
| return err |
| } |
| |
| if len(multiErr) > 0 { |
| fmt.Fprintf(stderr, "Failed to upload some rows:\n") |
| for _, rowErr := range multiErr { |
| for _, valErr := range rowErr.Errors { |
| fmt.Fprintf(stderr, "row %d: %s\n", rowErr.RowIndex, valErr) |
| } |
| } |
| return multiErr |
| } |
| return nil |
| } |
| |
| func overrideColumns(row map[string]bigquery.Value, columns map[string]bigquery.Value) map[string]bigquery.Value { |
| for k, v := range columns { |
| row[k] = v |
| } |
| return row |
| } |
| |
| func readInput(r io.Reader, insertIDBase string, jsonList bool, overrides map[string]bigquery.Value) (rows []*tableRow, err error) { |
| if jsonList { |
| var target []map[string]bigquery.Value |
| if err := json.NewDecoder(r).Decode(&target); err != nil { |
| return nil, err |
| } |
| rows = make([]*tableRow, len(target)) |
| for i, row := range target { |
| rows[i] = &tableRow{overrideColumns(row, overrides), fmt.Sprintf("%s:%d", insertIDBase, i)} |
| } |
| return |
| } |
| |
| buf := bufio.NewReaderSize(r, 32768) |
| |
| lineNo := 0 |
| for { |
| lineNo++ |
| |
| line, err := buf.ReadBytes('\n') |
| switch { |
| case err != nil && err != io.EOF: |
| return nil, err // a fatal error |
| case err == io.EOF && len(line) == 0: |
| return rows, nil // read past the last line |
| } |
| |
| if line = bytes.TrimSpace(line); len(line) != 0 { |
| row, err := parseRow(line, fmt.Sprintf("%s:%d", insertIDBase, len(rows))) |
| if err != nil { |
| return nil, fmt.Errorf("bad input line %d: %s", lineNo, err) |
| } |
| row.data = overrideColumns(row.data, overrides) |
| rows = append(rows, row) |
| } |
| } |
| } |
| |
| // tableRow implements bigquery.ValueSaver. |
| type tableRow struct { |
| data map[string]bigquery.Value |
| insertID string |
| } |
| |
| func parseRow(data []byte, insertID string) (*tableRow, error) { |
| row := make(map[string]bigquery.Value) |
| if err := json.Unmarshal(data, &row); err != nil { |
| return nil, fmt.Errorf("bad JSON - %s", err) |
| } |
| return &tableRow{row, insertID}, nil |
| } |
| |
| func (r *tableRow) Save() (map[string]bigquery.Value, string, error) { |
| return r.data, r.insertID, nil |
| } |