go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/bisection/bqexporter/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 bqexporter
    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/api/iterator"
    24  	"google.golang.org/protobuf/proto"
    25  
    26  	bqpb "go.chromium.org/luci/bisection/proto/bq"
    27  	"go.chromium.org/luci/bisection/util/bqutil"
    28  	"go.chromium.org/luci/common/errors"
    29  	"go.chromium.org/luci/gae/service/info"
    30  )
    31  
    32  // NewClient creates a new client for exporting test analyses
    33  // via the BigQuery Write API.
    34  // projectID is the project ID of the GCP project.
    35  func NewClient(ctx context.Context, projectID string) (s *Client, reterr error) {
    36  	if projectID == "" {
    37  		return nil, errors.New("GCP Project must be specified")
    38  	}
    39  
    40  	bqClient, err := bqutil.Client(ctx, projectID)
    41  	if err != nil {
    42  		return nil, errors.Annotate(err, "creating BQ client").Err()
    43  	}
    44  	defer func() {
    45  		if reterr != nil {
    46  			bqClient.Close()
    47  		}
    48  	}()
    49  
    50  	mwClient, err := bqutil.NewWriterClient(ctx, projectID)
    51  	if err != nil {
    52  		return nil, errors.Annotate(err, "create managed writer client").Err()
    53  	}
    54  	return &Client{
    55  		projectID: projectID,
    56  		bqClient:  bqClient,
    57  		mwClient:  mwClient,
    58  	}, nil
    59  }
    60  
    61  // Close releases resources held by the client.
    62  func (client *Client) Close() (reterr error) {
    63  	// Ensure both bqClient and mwClient Close() methods
    64  	// are called, even if one panics or fails.
    65  	defer func() {
    66  		err := client.mwClient.Close()
    67  		if reterr == nil {
    68  			reterr = err
    69  		}
    70  	}()
    71  	return client.bqClient.Close()
    72  }
    73  
    74  // Client provides methods to export test analyses to BigQuery
    75  // via the BigQuery Write API.
    76  type Client struct {
    77  	// projectID is the name of the GCP project that contains LUCI Bisection datasets.
    78  	projectID string
    79  	bqClient  *bigquery.Client
    80  	mwClient  *managedwriter.Client
    81  }
    82  
    83  func (client *Client) EnsureSchema(ctx context.Context) error {
    84  	table := client.bqClient.Dataset(bqutil.InternalDatasetID).Table(testFailureAnalysesTableName)
    85  	if err := schemaApplyer.EnsureTable(ctx, table, tableMetadata); err != nil {
    86  		return errors.Annotate(err, "ensuring test_analyses table").Err()
    87  	}
    88  	return nil
    89  }
    90  
    91  // Insert inserts the given rows in BigQuery.
    92  func (client *Client) Insert(ctx context.Context, rows []*bqpb.TestAnalysisRow) error {
    93  	if err := client.EnsureSchema(ctx); err != nil {
    94  		return errors.Annotate(err, "ensure schema").Err()
    95  	}
    96  	tableName := fmt.Sprintf("projects/%s/datasets/%s/tables/%s", client.projectID, bqutil.InternalDatasetID, testFailureAnalysesTableName)
    97  	writer := bqutil.NewWriter(client.mwClient, tableName, tableSchemaDescriptor)
    98  	payload := make([]proto.Message, len(rows))
    99  	for i, r := range rows {
   100  		payload[i] = r
   101  	}
   102  	// We use pending stream instead of default stream here because
   103  	// default stream does not offer exactly-once insert.
   104  	return writer.AppendRowsWithPendingStream(ctx, payload)
   105  }
   106  
   107  type TestFailureAnalysisRow struct {
   108  	// We only need analysis ID for now.
   109  	AnalysisID int64
   110  }
   111  
   112  // ReadTestFailureAnalysisRows returns the Test Failure analysis rows
   113  // in test_failure_analyses table that has created_time within the past 14 days.
   114  func (client *Client) ReadTestFailureAnalysisRows(ctx context.Context) ([]*TestFailureAnalysisRow, error) {
   115  	queryStm := fmt.Sprintf(`
   116  		SELECT DISTINCT
   117  			analysis_id as AnalysisID
   118  		FROM test_failure_analyses
   119  		WHERE created_time >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL %d DAY)
   120   	`, daysToLookBack)
   121  	q := client.bqClient.Query(queryStm)
   122  	q.DefaultDatasetID = bqutil.InternalDatasetID
   123  	q.DefaultProjectID = info.AppID(ctx)
   124  	job, err := q.Run(ctx)
   125  	if err != nil {
   126  		return nil, errors.Annotate(err, "querying test failure analyses").Err()
   127  	}
   128  	it, err := job.Read(ctx)
   129  	if err != nil {
   130  		return nil, err
   131  	}
   132  	rows := []*TestFailureAnalysisRow{}
   133  	for {
   134  		row := &TestFailureAnalysisRow{}
   135  		err := it.Next(row)
   136  		if err == iterator.Done {
   137  			break
   138  		}
   139  		if err != nil {
   140  			return nil, errors.Annotate(err, "obtain next test failure analysis row").Err()
   141  		}
   142  		rows = append(rows, row)
   143  	}
   144  	return rows, nil
   145  }