go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/analysis/internal/testverdicts/client.go (about)

     1  // Copyright 2023 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 testverdicts
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  
    21  	"cloud.google.com/go/bigquery"
    22  	"cloud.google.com/go/bigquery/storage/managedwriter"
    23  	"google.golang.org/protobuf/proto"
    24  
    25  	"go.chromium.org/luci/common/bq"
    26  	"go.chromium.org/luci/common/errors"
    27  
    28  	"go.chromium.org/luci/analysis/internal/bqutil"
    29  	bqpb "go.chromium.org/luci/analysis/proto/bq"
    30  )
    31  
    32  // NewClient creates a new client for exporting test verdicts
    33  // via the BigQuery Write API.
    34  func NewClient(ctx context.Context, projectID string) (s *Client, reterr error) {
    35  	if projectID == "" {
    36  		return nil, errors.New("GCP Project must be specified")
    37  	}
    38  
    39  	bqClient, err := bqutil.Client(ctx, projectID)
    40  	if err != nil {
    41  		return nil, errors.Annotate(err, "creating BQ client").Err()
    42  	}
    43  	defer func() {
    44  		if reterr != nil {
    45  			// This method failed for some reason, clean up the
    46  			// BigQuery client. Swallow any error returned by the Close()
    47  			// call.
    48  			bqClient.Close()
    49  		}
    50  	}()
    51  
    52  	mwClient, err := bqutil.NewWriterClient(ctx, projectID)
    53  	if err != nil {
    54  		return nil, errors.Annotate(err, "creating managed writer client").Err()
    55  	}
    56  	return &Client{
    57  		projectID: projectID,
    58  		bqClient:  bqClient,
    59  		mwClient:  mwClient,
    60  	}, nil
    61  }
    62  
    63  // Close releases resources held by the client.
    64  func (c *Client) Close() (reterr error) {
    65  	// Ensure both bqClient and mwClient Close() methods
    66  	// are called, even if one panics or fails.
    67  	defer func() {
    68  		err := c.mwClient.Close()
    69  		if reterr == nil {
    70  			reterr = err
    71  		}
    72  	}()
    73  	return c.bqClient.Close()
    74  }
    75  
    76  // Client provides methods to export test verdicts to BigQuery
    77  // via the BigQuery Write API.
    78  type Client struct {
    79  	// projectID is the name of the GCP project that contains LUCI Analysis
    80  	// BigQuery datasets.
    81  	projectID string
    82  	bqClient  *bigquery.Client
    83  	mwClient  *managedwriter.Client
    84  }
    85  
    86  // schemaApplier ensures BQ schema matches the row proto definitions.
    87  var schemaApplier = bq.NewSchemaApplyer(bq.RegisterSchemaApplyerCache(1))
    88  
    89  func (c *Client) ensureSchema(ctx context.Context) error {
    90  	// On new deployments, the internal dataset may have to be manually
    91  	// created.
    92  	table := c.bqClient.Dataset(bqutil.InternalDatasetID).Table(tableName)
    93  	if err := schemaApplier.EnsureTable(ctx, table, tableMetadata); err != nil {
    94  		return errors.Annotate(err, "ensuring test verdicts table").Err()
    95  	}
    96  	return nil
    97  }
    98  
    99  // Insert inserts the given rows in BigQuery.
   100  func (c *Client) Insert(ctx context.Context, rows []*bqpb.TestVerdictRow) error {
   101  	if err := c.ensureSchema(ctx); err != nil {
   102  		return errors.Annotate(err, "ensure schema").Err()
   103  	}
   104  	tableName := fmt.Sprintf("projects/%s/datasets/%s/tables/%s", c.projectID, bqutil.InternalDatasetID, tableName)
   105  	writer := bqutil.NewWriter(c.mwClient, tableName, tableSchemaDescriptor)
   106  	payload := make([]proto.Message, len(rows))
   107  	for i, r := range rows {
   108  		payload[i] = r
   109  	}
   110  	return writer.AppendRowsWithDefaultStream(ctx, payload)
   111  }