go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/cv/internal/common/bq/bq.go (about)

     1  // Copyright 2021 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 bq handles sending rows to BigQuery.
    16  package bq
    17  
    18  import (
    19  	"context"
    20  	"net/http"
    21  
    22  	"cloud.google.com/go/bigquery"
    23  
    24  	"google.golang.org/api/option"
    25  	"google.golang.org/protobuf/proto"
    26  
    27  	lucibq "go.chromium.org/luci/common/bq"
    28  	"go.chromium.org/luci/common/errors"
    29  	"go.chromium.org/luci/common/retry/transient"
    30  	"go.chromium.org/luci/server/auth"
    31  )
    32  
    33  // Row encapsulates destination and actual row to send.
    34  //
    35  // Exists to avoid confusion over multiple string arguments to SendRow.
    36  type Row struct {
    37  	// CloudProject allows sending rows to other projects.
    38  	//
    39  	// Optional. Defaults to the one in the scope of which this process is running
    40  	// (e.g. "luci-change-verifier-dev").
    41  	CloudProject string
    42  	Dataset      string
    43  	Table        string
    44  	// OperationID is used for de-duplication, but over just 1 minute window :(
    45  	OperationID string
    46  	Payload     proto.Message
    47  }
    48  
    49  type Client interface {
    50  	// SendRow appends a row to a BigQuery table synchronously.
    51  	SendRow(ctx context.Context, row Row) error
    52  }
    53  
    54  // NewProdClient creates new production client.
    55  //
    56  // The specified cloud project should be the one, in the scope of which this
    57  // code is running, e.g. "luci-change-verifier-dev".
    58  func NewProdClient(ctx context.Context, cloudProject string) (*prodClient, error) {
    59  	t, err := auth.GetRPCTransport(ctx, auth.AsSelf, auth.WithScopes(auth.CloudOAuthScopes...))
    60  	if err != nil {
    61  		return nil, err
    62  	}
    63  	b, err := bigquery.NewClient(ctx, cloudProject, option.WithHTTPClient(&http.Client{Transport: t}))
    64  	if err != nil {
    65  		return nil, errors.Annotate(err, "failed to create BQ client").Err()
    66  	}
    67  	return &prodClient{b}, nil
    68  }
    69  
    70  // prodClient implements a BigQuery Client for production.
    71  type prodClient struct {
    72  	b *bigquery.Client
    73  }
    74  
    75  // SendRow sends a row to a real BigQuery table.
    76  func (c *prodClient) SendRow(ctx context.Context, row Row) error {
    77  	var table *bigquery.Table
    78  	if row.CloudProject == "" {
    79  		table = c.b.Dataset(row.Dataset).Table(row.Table)
    80  	} else {
    81  		table = c.b.DatasetInProject(row.CloudProject, row.Dataset).Table(row.Table)
    82  	}
    83  	r := &lucibq.Row{
    84  		Message:  row.Payload,
    85  		InsertID: row.OperationID,
    86  	}
    87  	if err := table.Inserter().Put(ctx, r); err != nil {
    88  		if pme, _ := err.(bigquery.PutMultiError); len(pme) != 0 {
    89  			return errors.Annotate(err, "bad row").Err()
    90  		}
    91  		return errors.Annotate(err, "unknown error sending row").Tag(transient.Tag).Err()
    92  	}
    93  	return nil
    94  }