| // Copyright 2023 Google LLC |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // https://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| |
| package managedwriter |
| |
| import ( |
| "context" |
| "io" |
| "testing" |
| |
| "cloud.google.com/go/bigquery/storage/apiv1/storagepb" |
| "cloud.google.com/go/bigquery/storage/managedwriter/testdata" |
| "github.com/google/go-cmp/cmp" |
| "google.golang.org/protobuf/proto" |
| "google.golang.org/protobuf/reflect/protodesc" |
| "google.golang.org/protobuf/testing/protocmp" |
| "google.golang.org/protobuf/types/descriptorpb" |
| ) |
| |
| func TestSendOptimizer(t *testing.T) { |
| |
| exampleReq := &storagepb.AppendRowsRequest{ |
| Rows: &storagepb.AppendRowsRequest_ProtoRows{ |
| ProtoRows: &storagepb.AppendRowsRequest_ProtoData{ |
| Rows: &storagepb.ProtoRows{ |
| SerializedRows: [][]byte{[]byte("row_data")}, |
| }, |
| }, |
| }, |
| } |
| exampleStreamID := "foo" |
| exampleTraceID := "trace_id" |
| exampleReqFull := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest) |
| exampleReqFull.WriteStream = exampleStreamID |
| exampleReqFull.TraceId = buildTraceID(&streamSettings{TraceID: exampleTraceID}) |
| exampleDP := &descriptorpb.DescriptorProto{Name: proto.String("schema")} |
| exampleReqFull.GetProtoRows().WriterSchema = &storagepb.ProtoSchema{ |
| ProtoDescriptor: proto.Clone(exampleDP).(*descriptorpb.DescriptorProto), |
| } |
| |
| ctx := context.Background() |
| |
| var testCases = []struct { |
| description string |
| optimizer sendOptimizer |
| reqs []*pendingWrite |
| sendResults []error |
| wantReqs []*storagepb.AppendRowsRequest |
| }{ |
| { |
| description: "verbose-optimizer", |
| optimizer: &verboseOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmpl := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| return []*pendingWrite{ |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| } |
| }(), |
| sendResults: []error{ |
| nil, |
| io.EOF, |
| io.EOF, |
| }, |
| wantReqs: []*storagepb.AppendRowsRequest{ |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| }, |
| }, |
| { |
| description: "simplex no errors", |
| optimizer: &simplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmpl := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| return []*pendingWrite{ |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| } |
| }(), |
| sendResults: []error{ |
| nil, |
| nil, |
| nil, |
| }, |
| wantReqs: func() []*storagepb.AppendRowsRequest { |
| want := make([]*storagepb.AppendRowsRequest, 3) |
| // first has no redactions. |
| want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest) |
| req.GetProtoRows().WriterSchema = nil |
| req.TraceId = "" |
| req.WriteStream = "" |
| // second and third are optimized. |
| want[1] = req |
| want[2] = req |
| return want |
| }(), |
| }, |
| { |
| description: "simplex w/partial errors", |
| optimizer: &simplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmpl := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| return []*pendingWrite{ |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| } |
| }(), |
| sendResults: []error{ |
| nil, |
| io.EOF, |
| nil, |
| }, |
| wantReqs: func() []*storagepb.AppendRowsRequest { |
| want := make([]*storagepb.AppendRowsRequest, 3) |
| want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest) |
| req.GetProtoRows().WriterSchema = nil |
| req.TraceId = "" |
| req.WriteStream = "" |
| // second request is optimized |
| want[1] = req |
| // error causes third request to be full again. |
| want[2] = want[0] |
| return want |
| }(), |
| }, |
| { |
| description: "multiplex single all errors", |
| optimizer: &multiplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmpl := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| return []*pendingWrite{ |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| } |
| }(), |
| sendResults: []error{ |
| io.EOF, |
| io.EOF, |
| io.EOF, |
| }, |
| wantReqs: []*storagepb.AppendRowsRequest{ |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest), |
| }, |
| }, |
| { |
| description: "multiplex single no errors", |
| optimizer: &multiplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmpl := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| return []*pendingWrite{ |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| newPendingWrite(ctx, nil, proto.Clone(exampleReq).(*storagepb.AppendRowsRequest), tmpl, exampleStreamID, exampleTraceID), |
| } |
| }(), |
| sendResults: []error{ |
| nil, |
| nil, |
| nil, |
| }, |
| wantReqs: func() []*storagepb.AppendRowsRequest { |
| want := make([]*storagepb.AppendRowsRequest, 3) |
| want[0] = proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| req := proto.Clone(want[0]).(*storagepb.AppendRowsRequest) |
| req.GetProtoRows().WriterSchema = nil |
| req.TraceId = "" |
| want[1] = req |
| want[2] = req |
| return want |
| }(), |
| }, |
| { |
| description: "multiplex interleave", |
| optimizer: &multiplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmplA := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| tmplB := newVersionedTemplate().revise(reviseProtoSchema(protodesc.ToDescriptorProto((&testdata.AllSupportedTypes{}).ProtoReflect().Descriptor()))) |
| |
| reqA := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest) |
| reqA.WriteStream = "alpha" |
| |
| reqB := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest) |
| reqB.WriteStream = "beta" |
| |
| writes := make([]*pendingWrite, 10) |
| writes[0] = newPendingWrite(ctx, nil, reqA, tmplA, reqA.GetWriteStream(), exampleTraceID) |
| writes[1] = newPendingWrite(ctx, nil, reqA, tmplA, reqA.GetWriteStream(), exampleTraceID) |
| writes[2] = newPendingWrite(ctx, nil, reqB, tmplB, reqB.GetWriteStream(), exampleTraceID) |
| writes[3] = newPendingWrite(ctx, nil, reqA, tmplA, reqA.GetWriteStream(), exampleTraceID) |
| writes[4] = newPendingWrite(ctx, nil, reqB, tmplB, reqB.GetWriteStream(), exampleTraceID) |
| writes[5] = newPendingWrite(ctx, nil, reqB, tmplB, reqB.GetWriteStream(), exampleTraceID) |
| writes[6] = newPendingWrite(ctx, nil, reqB, tmplB, reqB.GetWriteStream(), exampleTraceID) |
| writes[7] = newPendingWrite(ctx, nil, reqB, tmplB, reqB.GetWriteStream(), exampleTraceID) |
| writes[8] = newPendingWrite(ctx, nil, reqA, tmplA, reqA.GetWriteStream(), exampleTraceID) |
| writes[9] = newPendingWrite(ctx, nil, reqA, tmplA, reqA.GetWriteStream(), exampleTraceID) |
| |
| return writes |
| }(), |
| sendResults: []error{ |
| nil, |
| nil, |
| nil, |
| nil, |
| nil, |
| io.EOF, |
| nil, |
| nil, |
| nil, |
| io.EOF, |
| }, |
| wantReqs: func() []*storagepb.AppendRowsRequest { |
| want := make([]*storagepb.AppendRowsRequest, 10) |
| |
| wantReqAFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| wantReqAFull.WriteStream = "alpha" |
| |
| wantReqANoTrace := proto.Clone(wantReqAFull).(*storagepb.AppendRowsRequest) |
| wantReqANoTrace.TraceId = "" |
| |
| wantReqAOpt := proto.Clone(wantReqAFull).(*storagepb.AppendRowsRequest) |
| wantReqAOpt.GetProtoRows().WriterSchema = nil |
| wantReqAOpt.TraceId = "" |
| |
| wantReqBFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| wantReqBFull.WriteStream = "beta" |
| wantReqBFull.GetProtoRows().GetWriterSchema().ProtoDescriptor = protodesc.ToDescriptorProto((&testdata.AllSupportedTypes{}).ProtoReflect().Descriptor()) |
| |
| wantReqBNoTrace := proto.Clone(wantReqBFull).(*storagepb.AppendRowsRequest) |
| wantReqBNoTrace.TraceId = "" |
| |
| wantReqBOpt := proto.Clone(wantReqBFull).(*storagepb.AppendRowsRequest) |
| wantReqBOpt.GetProtoRows().WriterSchema = nil |
| wantReqBOpt.TraceId = "" |
| |
| want[0] = wantReqAFull |
| want[1] = wantReqAOpt |
| want[2] = wantReqBNoTrace |
| want[3] = wantReqANoTrace |
| want[4] = wantReqBNoTrace |
| want[5] = wantReqBOpt |
| want[6] = wantReqBFull |
| want[7] = wantReqBOpt |
| want[8] = wantReqANoTrace |
| want[9] = wantReqAOpt |
| |
| return want |
| }(), |
| }, |
| { |
| description: "multiplex w/evolution", |
| optimizer: &multiplexOptimizer{}, |
| reqs: func() []*pendingWrite { |
| tmplOld := newVersionedTemplate().revise(reviseProtoSchema(exampleDP)) |
| tmplNew := tmplOld.revise(reviseProtoSchema(&descriptorpb.DescriptorProto{Name: proto.String("new")})) |
| |
| example := proto.Clone(exampleReq).(*storagepb.AppendRowsRequest) |
| |
| writes := make([]*pendingWrite, 4) |
| writes[0] = newPendingWrite(ctx, nil, example, tmplOld, exampleStreamID, exampleTraceID) |
| writes[1] = newPendingWrite(ctx, nil, example, tmplOld, exampleStreamID, exampleTraceID) |
| writes[2] = newPendingWrite(ctx, nil, example, tmplNew, exampleStreamID, exampleTraceID) |
| writes[3] = newPendingWrite(ctx, nil, example, tmplNew, exampleStreamID, exampleTraceID) |
| |
| return writes |
| }(), |
| sendResults: []error{ |
| nil, |
| nil, |
| nil, |
| nil, |
| }, |
| wantReqs: func() []*storagepb.AppendRowsRequest { |
| want := make([]*storagepb.AppendRowsRequest, 4) |
| |
| wantBaseReqFull := proto.Clone(exampleReqFull).(*storagepb.AppendRowsRequest) |
| |
| wantBaseReqOpt := proto.Clone(wantBaseReqFull).(*storagepb.AppendRowsRequest) |
| wantBaseReqOpt.TraceId = "" |
| wantBaseReqOpt.GetProtoRows().WriterSchema = nil |
| |
| wantEvolved := proto.Clone(wantBaseReqOpt).(*storagepb.AppendRowsRequest) |
| wantEvolved.GetProtoRows().WriterSchema = &storagepb.ProtoSchema{ |
| ProtoDescriptor: &descriptorpb.DescriptorProto{Name: proto.String("new")}, |
| } |
| |
| want[0] = wantBaseReqFull |
| want[1] = wantBaseReqOpt |
| want[2] = wantEvolved |
| want[3] = wantBaseReqOpt |
| return want |
| }(), |
| }, |
| } |
| |
| for _, tc := range testCases { |
| testARC := &testAppendRowsClient{} |
| testARC.sendF = func(req *storagepb.AppendRowsRequest) error { |
| testARC.requests = append(testARC.requests, proto.Clone(req).(*storagepb.AppendRowsRequest)) |
| respErr := tc.sendResults[0] |
| tc.sendResults = tc.sendResults[1:] |
| return respErr |
| } |
| |
| for _, req := range tc.reqs { |
| err := tc.optimizer.optimizeSend(testARC, req) |
| if err != nil { |
| tc.optimizer.signalReset() |
| } |
| } |
| // now, compare. |
| for k, wr := range tc.wantReqs { |
| if diff := cmp.Diff(testARC.requests[k], wr, protocmp.Transform()); diff != "" { |
| t.Errorf("%s (req %d) mismatch: -got, +want:\n%s", tc.description, k, diff) |
| } |
| } |
| } |
| } |
| |
| func TestVersionedTemplate(t *testing.T) { |
| testCases := []struct { |
| desc string |
| inputTmpl *storagepb.AppendRowsRequest |
| changes []templateRevisionF |
| wantCompatible bool |
| }{ |
| { |
| desc: "nil template", |
| wantCompatible: true, |
| }, |
| { |
| desc: "no changes", |
| inputTmpl: &storagepb.AppendRowsRequest{}, |
| wantCompatible: true, |
| }, |
| { |
| desc: "empty schema", |
| inputTmpl: &storagepb.AppendRowsRequest{}, |
| changes: []templateRevisionF{ |
| reviseProtoSchema(nil), |
| }, |
| wantCompatible: false, |
| }, |
| { |
| desc: "same default mvi", |
| inputTmpl: &storagepb.AppendRowsRequest{ |
| DefaultMissingValueInterpretation: storagepb.AppendRowsRequest_NULL_VALUE, |
| }, |
| changes: []templateRevisionF{ |
| reviseDefaultMissingValueInterpretation(storagepb.AppendRowsRequest_NULL_VALUE), |
| }, |
| wantCompatible: true, |
| }, |
| { |
| desc: "differing default mvi", |
| inputTmpl: &storagepb.AppendRowsRequest{ |
| DefaultMissingValueInterpretation: storagepb.AppendRowsRequest_NULL_VALUE, |
| }, |
| changes: []templateRevisionF{ |
| reviseDefaultMissingValueInterpretation(storagepb.AppendRowsRequest_DEFAULT_VALUE), |
| }, |
| wantCompatible: false, |
| }, |
| } |
| |
| for _, tc := range testCases { |
| orig := newVersionedTemplate() |
| orig.tmpl = tc.inputTmpl |
| orig.computeHash() |
| |
| rev := orig.revise(tc.changes...) |
| if orig.Compatible(rev) != rev.Compatible(orig) { |
| t.Errorf("case %q: inconsistent compatibility, orig %t rev %t", tc.desc, orig.Compatible(rev), rev.Compatible(orig)) |
| } |
| if got := orig.Compatible(rev); tc.wantCompatible != got { |
| t.Errorf("case %q: Compatible mismatch, got %t want %t", tc.desc, got, tc.wantCompatible) |
| } |
| } |
| } |
