go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/analysis/internal/changepoints/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  	"cloud.google.com/go/bigquery"
    19  	"cloud.google.com/go/bigquery/storage/managedwriter/adapt"
    20  	"github.com/golang/protobuf/descriptor"
    21  	desc "github.com/golang/protobuf/protoc-gen-go/descriptor"
    22  	"google.golang.org/protobuf/types/descriptorpb"
    23  
    24  	"go.chromium.org/luci/common/bq"
    25  
    26  	"go.chromium.org/luci/analysis/internal/bqutil"
    27  	bqpb "go.chromium.org/luci/analysis/proto/bq"
    28  	pb "go.chromium.org/luci/analysis/proto/v1"
    29  )
    30  
    31  // We stream test variant branch updates to this table.
    32  const updatesTableName = "test_variant_segment_updates"
    33  
    34  // This is the main table for the test variant branches.
    35  // Periodically we merge from the test_variant_segment_updates to this table.
    36  const stableTableName = "test_variant_segments"
    37  
    38  // schemaApplyer ensures BQ schema matches the row proto definitions.
    39  var schemaApplyer = bq.NewSchemaApplyer(bq.RegisterSchemaApplyerCache(50))
    40  
    41  const rowMessage = "luci.analysis.bq.TestVariantBranchRow"
    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  		RangePartitioning: &bigquery.RangePartitioning{
    61  			// has_recent_unexpected_results has only 2 possible values: 0 and 1.
    62  			Field: "has_recent_unexpected_results",
    63  			Range: &bigquery.RangePartitioningRange{
    64  				Start: 0,
    65  				// End is exclusive.
    66  				End: 2,
    67  				// The width of each interval range.
    68  				Interval: 1,
    69  			},
    70  		},
    71  		Clustering: &bigquery.Clustering{
    72  			Fields: []string{"project", "test_id"},
    73  		},
    74  		// Relax ensures no fields are marked "required".
    75  		Schema: schema.Relax(),
    76  	}
    77  }
    78  
    79  func generateRowSchema() (schema bigquery.Schema, err error) {
    80  	fd, _ := descriptor.MessageDescriptorProto(&bqpb.TestVariantBranchRow{})
    81  	// We also need to get FileDescriptorProto for SourceRef and GitilesRef
    82  	// because they are defined in different files.
    83  	fdsr, _ := descriptor.MessageDescriptorProto(&pb.SourceRef{})
    84  	fdgr, _ := descriptor.MessageDescriptorProto(&pb.GitilesRef{})
    85  	fdset := &desc.FileDescriptorSet{File: []*desc.FileDescriptorProto{fd, fdsr, fdgr}}
    86  	return bqutil.GenerateSchema(fdset, rowMessage)
    87  }
    88  
    89  func generateRowSchemaDescriptor() (*desc.DescriptorProto, error) {
    90  	m := &bqpb.TestVariantBranchRow{}
    91  	descriptorProto, err := adapt.NormalizeDescriptor(m.ProtoReflect().Descriptor())
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  	return descriptorProto, nil
    96  }