go.fuchsia.dev/infra@v0.0.0-20240507153436-9b593402251b/cmd/update_test_durations/query.go (about)

     1  // Copyright 2022 The Fuchsia Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	_ "embed"
    10  	"errors"
    11  	"time"
    12  
    13  	"cloud.google.com/go/bigquery"
    14  	"google.golang.org/api/iterator"
    15  )
    16  
    17  const (
    18  	cloudProject = "fuchsia-infra"
    19  
    20  	// Timeout for BigQuery queries.
    21  	queryTimeout = 10 * time.Minute
    22  )
    23  
    24  //go:embed query.sql
    25  var queryTemplate string
    26  
    27  // queryLatestTestDurations runs a BigQuery query to obtain recent test duration
    28  // data.
    29  func queryLatestTestDurations(ctx context.Context, luciProject string, dataWindowDays int) ([]test, error) {
    30  	ctx, cancel := context.WithTimeout(ctx, queryTimeout)
    31  	defer cancel()
    32  
    33  	client, err := bigquery.NewClient(ctx, cloudProject)
    34  	if err != nil {
    35  		return nil, err
    36  	}
    37  	defer client.Close()
    38  
    39  	query := client.Query(queryTemplate)
    40  	query.Parameters = append(query.Parameters, []bigquery.QueryParameter{
    41  		{Name: "data_window_days", Value: dataWindowDays},
    42  		{Name: "project", Value: luciProject},
    43  	}...)
    44  
    45  	iter, err := query.Read(ctx)
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	var rows []test
    51  	for {
    52  		var row test
    53  		err := iter.Next(&row)
    54  		if err == iterator.Done {
    55  			break
    56  		} else if err != nil {
    57  			return nil, err
    58  		}
    59  		rows = append(rows, row)
    60  	}
    61  
    62  	if len(rows) == 0 {
    63  		return nil, errors.New("query returned zero rows")
    64  	}
    65  
    66  	return rows, nil
    67  }