feat(bigquery/storage/managedwriter): support variadic appends (#5102)

* feat(bigquery/storage/managedwriter): support variadic AppendOption

Author: shollyman

Diff for: bigquery/storage/managedwriter/appendresult.go

@@ -19,6 +19,7 @@ import (
 
 	storagepb "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/descriptorpb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
@@ -66,7 +67,9 @@ func (ar *AppendResult) GetResult(ctx context.Context) (int64, error) {
 // append request.
 type pendingWrite struct {
 	request *storagepb.AppendRowsRequest
-	result  *AppendResult
+	// for schema evolution cases, accept a new schema
+	newSchema *descriptorpb.DescriptorProto
+	result    *AppendResult
 
 	// this is used by the flow controller.
 	reqSize int
@@ -77,7 +80,7 @@ type pendingWrite struct {
 // that in the future, we may want to allow row batching to be managed by
 // the server (e.g. for default/COMMITTED streams).  For BUFFERED/PENDING
 // streams, this should be managed by the user.
-func newPendingWrite(appends [][]byte, offset int64) *pendingWrite {
+func newPendingWrite(appends [][]byte) *pendingWrite {
 	pw := &pendingWrite{
 		request: &storagepb.AppendRowsRequest{
 			Rows: &storagepb.AppendRowsRequest_ProtoRows{
@@ -90,9 +93,6 @@ func newPendingWrite(appends [][]byte, offset int64) *pendingWrite {
 		},
 		result: newAppendResult(appends),
 	}
-	if offset > 0 {
-		pw.request.Offset = &wrapperspb.Int64Value{Value: offset}
-	}
 	// We compute the size now for flow controller purposes, though
 	// the actual request size may be slightly larger (e.g. the first
 	// request in a new stream bears schema and stream id).
@@ -114,3 +114,23 @@ func (pw *pendingWrite) markDone(startOffset int64, err error, fc *flowControlle
 		fc.release(pw.reqSize)
 	}
 }
+
+// AppendOption are options that can be passed when appending data with a managed stream instance.
+type AppendOption func(*pendingWrite)
+
+// UpdateSchemaDescriptor is used to update the descriptor message schema associated
+// with a given stream.
+func UpdateSchemaDescriptor(schema *descriptorpb.DescriptorProto) AppendOption {
+	return func(pw *pendingWrite) {
+		pw.newSchema = schema
+	}
+}
+
+// WithOffset sets an explicit offset value for this append request.
+func WithOffset(offset int64) AppendOption {
+	return func(pw *pendingWrite) {
+		pw.request.Offset = &wrapperspb.Int64Value{
+			Value: offset,
+		}
+	}
+}

Diff for: bigquery/storage/managedwriter/appendresult_test.go

@@ -43,38 +43,33 @@ func TestPendingWrite(t *testing.T) {
 		[]byte("row3"),
 	}
 
-	var wantOffset int64 = 99
-
-	// first, verify no offset behavior
-	pending := newPendingWrite(wantRowData, NoStreamOffset)
+	// verify no offset behavior
+	pending := newPendingWrite(wantRowData)
 	if pending.request.GetOffset() != nil {
 		t.Errorf("request should have no offset, but is present: %q", pending.request.GetOffset().GetValue())
 	}
-	pending.markDone(NoStreamOffset, nil, nil)
-	if pending.result.offset != NoStreamOffset {
-		t.Errorf("mismatch on completed AppendResult without offset: got %d want %d", pending.result.offset, NoStreamOffset)
-	}
-	if pending.result.err != nil {
-		t.Errorf("mismatch in error on AppendResult, got %v want nil", pending.result.err)
-	}
-
-	// now, verify behavior with a valid offset
-	pending = newPendingWrite(wantRowData, 99)
-	if pending.request.GetOffset() == nil {
-		t.Errorf("offset not set, should be %d", wantOffset)
-	}
-	if gotOffset := pending.request.GetOffset().GetValue(); gotOffset != wantOffset {
-		t.Errorf("offset mismatch, got %d want %d", gotOffset, wantOffset)
-	}
 
-	// check request shape
 	gotRowCount := len(pending.request.GetProtoRows().GetRows().GetSerializedRows())
 	if gotRowCount != len(wantRowData) {
 		t.Errorf("pendingWrite request mismatch, got %d rows, want %d rows", gotRowCount, len(wantRowData))
 	}
 
-	// verify AppendResult
+	// Verify request is not acknowledged.
+	select {
+	case <-pending.result.Ready():
+		t.Errorf("got Ready() on incomplete AppendResult")
+	case <-time.After(100 * time.Millisecond):
+
+	}
 
+	// Mark completed, verify result.
+	pending.markDone(NoStreamOffset, nil, nil)
+	if pending.result.offset != NoStreamOffset {
+		t.Errorf("mismatch on completed AppendResult without offset: got %d want %d", pending.result.offset, NoStreamOffset)
+	}
+	if pending.result.err != nil {
+		t.Errorf("mismatch in error on AppendResult, got %v want nil", pending.result.err)
+	}
 	gotData := pending.result.rowData
 	if len(gotData) != len(wantRowData) {
 		t.Errorf("length mismatch on appendresult, got %d, want %d", len(gotData), len(wantRowData))
@@ -84,15 +79,24 @@ func TestPendingWrite(t *testing.T) {
 			t.Errorf("row %d mismatch in data: got %q want %q", i, gotData[i], wantRowData[i])
 		}
 	}
-	select {
-	case <-pending.result.Ready():
-		t.Errorf("got Ready() on incomplete AppendResult")
-	case <-time.After(100 * time.Millisecond):
 
-	}
+	// Create new write to verify error result.
+	pending = newPendingWrite(wantRowData)
 
-	// verify completion behavior
+	// Manually invoke option to apply offset to request.
+	// This would normally be appied as part of the AppendRows() method on the managed stream.
 	reportedOffset := int64(101)
+	f := WithOffset(reportedOffset)
+	f(pending)
+
+	if pending.request.GetOffset() == nil {
+		t.Errorf("expected offset, got none")
+	}
+	if pending.request.GetOffset().GetValue() != reportedOffset {
+		t.Errorf("offset mismatch, got %d wanted %d", pending.request.GetOffset().GetValue(), reportedOffset)
+	}
+
+	// Verify completion behavior with an error.
 	wantErr := fmt.Errorf("foo")
 	pending.markDone(reportedOffset, wantErr, nil)
 

Diff for: bigquery/storage/managedwriter/integration_test.go

@@ -144,6 +144,10 @@ func TestIntegration_ManagedWriter(t *testing.T) {
 			t.Parallel()
 			testPendingStream(ctx, t, mwClient, bqClient, dataset)
 		})
+		t.Run("SchemaEvolution", func(t *testing.T) {
+			t.Parallel()
+			testSchemaEvolution(ctx, t, mwClient, bqClient, dataset)
+		})
 		t.Run("Instrumentation", func(t *testing.T) {
 			// Don't run this in parallel, we only want to collect stats from this subtest.
 			testInstrumentation(ctx, t, mwClient, bqClient, dataset)
@@ -181,7 +185,7 @@ func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
 			t.Errorf("failed to marshal message %d: %v", k, err)
 		}
 		data := [][]byte{b}
-		result, err = ms.AppendRows(ctx, data, NoStreamOffset)
+		result, err = ms.AppendRows(ctx, data)
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
@@ -201,7 +205,7 @@ func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
 		}
 		data = append(data, b)
 	}
-	result, err = ms.AppendRows(ctx, data, NoStreamOffset)
+	result, err = ms.AppendRows(ctx, data)
 	if err != nil {
 		t.Errorf("grouped-row append failed: %v", err)
 	}
@@ -256,7 +260,7 @@ func testDefaultStreamDynamicJSON(ctx context.Context, t *testing.T, mwClient *C
 		if err != nil {
 			t.Fatalf("failed to marshal proto bytes for row %d: %v", k, err)
 		}
-		result, err = ms.AppendRows(ctx, [][]byte{b}, NoStreamOffset)
+		result, err = ms.AppendRows(ctx, [][]byte{b})
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
@@ -305,7 +309,7 @@ func testBufferedStream(ctx context.Context, t *testing.T, mwClient *Client, bqC
 			t.Errorf("failed to marshal message %d: %v", k, err)
 		}
 		data := [][]byte{b}
-		results, err := ms.AppendRows(ctx, data, NoStreamOffset)
+		results, err := ms.AppendRows(ctx, data)
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
@@ -358,7 +362,7 @@ func testCommittedStream(ctx context.Context, t *testing.T, mwClient *Client, bq
 			t.Errorf("failed to marshal message %d: %v", k, err)
 		}
 		data := [][]byte{b}
-		result, err = ms.AppendRows(ctx, data, NoStreamOffset)
+		result, err = ms.AppendRows(ctx, data, WithOffset(int64(k)))
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
@@ -397,12 +401,19 @@ func testPendingStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
 			t.Errorf("failed to marshal message %d: %v", k, err)
 		}
 		data := [][]byte{b}
-		result, err = ms.AppendRows(ctx, data, NoStreamOffset)
+		result, err = ms.AppendRows(ctx, data, WithOffset(int64(k)))
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
+		// be explicit about waiting/checking each response.
+		off, err := result.GetResult(ctx)
+		if err != nil {
+			t.Errorf("response %d error: %v", k, err)
+		}
+		if off != int64(k) {
+			t.Errorf("offset mismatch, got %d want %d", off, k)
+		}
 	}
-	result.Ready()
 	wantRows := int64(len(testSimpleData))
 
 	// Mark stream complete.
@@ -468,7 +479,7 @@ func testInstrumentation(ctx context.Context, t *testing.T, mwClient *Client, bq
 			t.Errorf("failed to marshal message %d: %v", k, err)
 		}
 		data := [][]byte{b}
-		result, err = ms.AppendRows(ctx, data, NoStreamOffset)
+		result, err = ms.AppendRows(ctx, data)
 		if err != nil {
 			t.Errorf("single-row append %d failed: %v", k, err)
 		}
@@ -513,6 +524,85 @@ func testInstrumentation(ctx context.Context, t *testing.T, mwClient *Client, bq
 	}
 }
 
+func testSchemaEvolution(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {
+	testTable := dataset.Table(tableIDs.New())
+	if err := testTable.Create(ctx, &bigquery.TableMetadata{Schema: testdata.SimpleMessageSchema}); err != nil {
+		t.Fatalf("failed to create test table %s: %v", testTable.FullyQualifiedName(), err)
+	}
+
+	m := &testdata.SimpleMessageProto2{}
+	descriptorProto := protodesc.ToDescriptorProto(m.ProtoReflect().Descriptor())
+
+	// setup a new stream.
+	ms, err := mwClient.NewManagedStream(ctx,
+		WithDestinationTable(fmt.Sprintf("projects/%s/datasets/%s/tables/%s", testTable.ProjectID, testTable.DatasetID, testTable.TableID)),
+		WithType(CommittedStream),
+		WithSchemaDescriptor(descriptorProto),
+	)
+	if err != nil {
+		t.Fatalf("NewManagedStream: %v", err)
+	}
+	validateTableConstraints(ctx, t, bqClient, testTable, "before send",
+		withExactRowCount(0))
+
+	var result *AppendResult
+	for k, mesg := range testSimpleData {
+		b, err := proto.Marshal(mesg)
+		if err != nil {
+			t.Errorf("failed to marshal message %d: %v", k, err)
+		}
+		data := [][]byte{b}
+		result, err = ms.AppendRows(ctx, data)
+		if err != nil {
+			t.Errorf("single-row append %d failed: %v", k, err)
+		}
+	}
+	// wait for the result to indicate ready, then validate.
+	_, err = result.GetResult(ctx)
+	if err != nil {
+		t.Errorf("error on append: %v", err)
+	}
+
+	validateTableConstraints(ctx, t, bqClient, testTable, "after send",
+		withExactRowCount(int64(len(testSimpleData))))
+
+	// Now, evolve the underlying table schema.
+	_, err = testTable.Update(ctx, bigquery.TableMetadataToUpdate{Schema: testdata.SimpleMessageEvolvedSchema}, "")
+	if err != nil {
+		t.Errorf("failed to evolve table schema: %v", err)
+	}
+
+	// TODO: we need a more elegant mechanism for detecting when the backend has registered the schema change.
+	//       In the continuous case, we'd get it from the response, but the change-and-wait case needs something more.
+	time.Sleep(6 * time.Second)
+
+	// ready descriptor, send an additional append
+	m2 := &testdata.SimpleMessageEvolvedProto2{
+		Name:  proto.String("evolved"),
+		Value: proto.Int64(180),
+		Other: proto.String("hello evolution"),
+	}
+	descriptorProto = protodesc.ToDescriptorProto(m2.ProtoReflect().Descriptor())
+	b, err := proto.Marshal(m2)
+	if err != nil {
+		t.Errorf("failed to marshal evolved message: %v", err)
+	}
+	result, err = ms.AppendRows(ctx, [][]byte{b}, UpdateSchemaDescriptor(descriptorProto))
+	if err != nil {
+		t.Errorf("failed evolved append: %v", err)
+	}
+	_, err = result.GetResult(ctx)
+	if err != nil {
+		t.Errorf("error on evolved append: %v", err)
+	}
+
+	validateTableConstraints(ctx, t, bqClient, testTable, "after send",
+		withExactRowCount(int64(len(testSimpleData)+1)),
+		withNullCount("name", 0),
+		withNonNullCount("other", 1),
+	)
+}
+
 func TestIntegration_DetectProjectID(t *testing.T) {
 	ctx := context.Background()
 	testCreds := testutil.Credentials(ctx)
@@ -613,7 +703,7 @@ func testProtoNormalization(ctx context.Context, t *testing.T, mwClient *Client,
 	if err != nil {
 		t.Fatalf("NewManagedStream: %v", err)
 	}
-	result, err := ms.AppendRows(ctx, [][]byte{sampleRow}, NoStreamOffset)
+	result, err := ms.AppendRows(ctx, [][]byte{sampleRow})
 	if err != nil {
 		t.Errorf("append failed: %v", err)
 	}