go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/bisection/bqexporter/schema.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  	"time"
    19  
    20  	"cloud.google.com/go/bigquery"
    21  	"cloud.google.com/go/bigquery/storage/managedwriter/adapt"
    22  	"github.com/golang/protobuf/descriptor"
    23  	desc "github.com/golang/protobuf/protoc-gen-go/descriptor"
    24  	"google.golang.org/protobuf/types/descriptorpb"
    25  
    26  	bqpb "go.chromium.org/luci/bisection/proto/bq"
    27  	pb "go.chromium.org/luci/bisection/proto/v1"
    28  	"go.chromium.org/luci/bisection/util/bqutil"
    29  	bbpb "go.chromium.org/luci/buildbucket/proto"
    30  	"go.chromium.org/luci/common/bq"
    31  )
    32  
    33  // The table containing test analyses.
    34  const testFailureAnalysesTableName = "test_failure_analyses"
    35  
    36  const partitionExpirationTime = 540 * 24 * time.Hour // 540 days (~1.5 years).
    37  
    38  // schemaApplyer ensures BQ schema matches the row proto definitions.
    39  var schemaApplyer = bq.NewSchemaApplyer(bq.RegisterSchemaApplyerCache(5))
    40  
    41  const rowMessage = "luci.bisection.proto.bq.TestAnalysisRow"
    42  
    43  var tableMetadata *bigquery.TableMetadata
    44  
    45  // tableSchemaDescriptor is a self-contained DescriptorProto for describing
    46  // row protocol buffers sent to the BigQuery Write API.
    47  var tableSchemaDescriptor *descriptorpb.DescriptorProto
    48  
    49  func init() {
    50  	var err error
    51  	var schema bigquery.Schema
    52  	if schema, err = generateRowSchema(); err != nil {
    53  		panic(err)
    54  	}
    55  	if tableSchemaDescriptor, err = generateRowSchemaDescriptor(); err != nil {
    56  		panic(err)
    57  	}
    58  
    59  	tableMetadata = &bigquery.TableMetadata{
    60  		TimePartitioning: &bigquery.TimePartitioning{
    61  			Type:       bigquery.DayPartitioningType,
    62  			Expiration: partitionExpirationTime,
    63  			Field:      "created_time",
    64  		},
    65  		// Relax ensures no fields are marked "required".
    66  		Schema: schema.Relax(),
    67  	}
    68  }
    69  
    70  func generateRowSchema() (schema bigquery.Schema, err error) {
    71  	fd, _ := descriptor.MessageDescriptorProto(&bqpb.TestAnalysisRow{})
    72  	// We also need to get FileDescriptorProto for:
    73  	// - TestCulprit
    74  	// - BuilderID
    75  	// - TestFailure
    76  	// - GitilesCommit
    77  	// - TestNthSectionAnalysisResult
    78  	// - CulpritAction
    79  	// - Variant
    80  	// - RegressionRange
    81  	// because they are defined in different files.
    82  	fdtc, _ := descriptor.MessageDescriptorProto(&pb.TestCulprit{})
    83  	fdbid, _ := descriptor.MessageDescriptorProto(&bbpb.BuilderID{})
    84  	fdtf, _ := descriptor.MessageDescriptorProto(&pb.TestFailure{})
    85  	fdgc, _ := descriptor.MessageDescriptorProto(&bbpb.GitilesCommit{})
    86  	fdtns, _ := descriptor.MessageDescriptorProto(&pb.TestNthSectionAnalysisResult{})
    87  	fdca, _ := descriptor.MessageDescriptorProto(&pb.CulpritAction{})
    88  	fdvr, _ := descriptor.MessageDescriptorProto(&pb.Variant{})
    89  	fdrr, _ := descriptor.MessageDescriptorProto(&pb.RegressionRange{})
    90  	fdset := &desc.FileDescriptorSet{File: []*desc.FileDescriptorProto{fd, fdtc, fdbid, fdtf, fdgc, fdtns, fdca, fdvr, fdrr}}
    91  	return bqutil.GenerateSchema(fdset, rowMessage)
    92  }
    93  
    94  func generateRowSchemaDescriptor() (*desc.DescriptorProto, error) {
    95  	m := &bqpb.TestAnalysisRow{}
    96  	descriptorProto, err := adapt.NormalizeDescriptor(m.ProtoReflect().Descriptor())
    97  	if err != nil {
    98  		return nil, err
    99  	}
   100  	return descriptorProto, nil
   101  }