feat(bigquery/storage/managedwriter): refactor AppendResponse (#6402)

* feat(bigquery/storage/managedwriter): refactor AppendResponse

The potential fields exposed within an AppendResponse has grown as
the API has evolved.  This PR refactors AppendResult to use a
retained reference of the response for servicing requests.

This allows the logic for processing the response to be centralized
a bit more within the AppendResult.  We also introduce a new
FullResponse() on the AppendResult which returns the full
AppendRowsResponse if present.

bigquery/storage/managedwriter/appendresult.go

@@ -18,7 +18,9 @@ import (
 	"context"
 	"fmt"
 
+	"github.com/googleapis/gax-go/v2/apierror"
 	storagepb "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
+	grpcstatus "google.golang.org/grpc/status"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/descriptorpb"
 )
@@ -34,14 +36,11 @@ type AppendResult struct {
 
 	ready chan struct{}
 
-	// if the encapsulating append failed, this will retain a reference to the error.
+	// if the append failed without a response, this will retain a reference to the error.
 	err error
 
-	// the stream offset
-	offset int64
-
-	// retains the updated schema from backend response.  Used for schema change notification.
-	updatedSchema *storagepb.TableSchema
+	// retains the original response.
+	response *storagepb.AppendRowsResponse
 }
 
 func newAppendResult(data [][]byte) *AppendResult {
@@ -55,25 +54,95 @@ func newAppendResult(data [][]byte) *AppendResult {
 // which may be a successful append or an error.
 func (ar *AppendResult) Ready() <-chan struct{} { return ar.ready }
 
-// GetResult returns the optional offset of this row, or the associated
-// error.  It blocks until the result is ready.
+// GetResult returns the optional offset of this row, as well as any error encountered while
+// processing the append.
+//
+// This call blocks until the result is ready, or context is no longer valid.
 func (ar *AppendResult) GetResult(ctx context.Context) (int64, error) {
 	select {
 	case <-ctx.Done():
-		return 0, ctx.Err()
+		return NoStreamOffset, ctx.Err()
+	case <-ar.Ready():
+		full, err := ar.FullResponse(ctx)
+		offset := NoStreamOffset
+		if full != nil {
+			if result := full.GetAppendResult(); result != nil {
+				if off := result.GetOffset(); off != nil {
+					offset = off.GetValue()
+				}
+			}
+		}
+		return offset, err
+	}
+}
+
+// FullResponse returns the full content of the AppendRowsResponse, and any error encountered while
+// processing the append.
+//
+// The AppendRowResponse may contain an embedded error.  An embedded error in the response will be
+// converted and returned as the error response, so this method may return both the
+// AppendRowsResponse and an error.
+//
+// This call blocks until the result is ready, or context is no longer valid.
+func (ar *AppendResult) FullResponse(ctx context.Context) (*storagepb.AppendRowsResponse, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
 	case <-ar.Ready():
-		return ar.offset, ar.err
+		var err error
+		if ar.err != nil {
+			err = ar.err
+		} else {
+			if ar.response != nil {
+				if status := ar.response.GetError(); status != nil {
+					statusErr := grpcstatus.ErrorProto(status)
+					// Provide an APIError if possible.
+					if apiErr, ok := apierror.FromError(statusErr); ok {
+						err = apiErr
+					} else {
+						err = statusErr
+					}
+				}
+			}
+		}
+		if ar.response != nil {
+			return proto.Clone(ar.response).(*storagepb.AppendRowsResponse), err
+		}
+		return nil, err
+	}
+}
+
+func (ar *AppendResult) offset(ctx context.Context) int64 {
+	select {
+	case <-ctx.Done():
+		return NoStreamOffset
+	case <-ar.Ready():
+		if ar.response != nil {
+			if result := ar.response.GetAppendResult(); result != nil {
+				if off := result.GetOffset(); off != nil {
+					return off.GetValue()
+				}
+			}
+		}
+		return NoStreamOffset
 	}
 }
 
 // UpdatedSchema returns the updated schema for a table if supplied by the backend as part
-// of the append response.  It blocks until the result is ready.
+// of the append response.
+//
+// This call blocks until the result is ready, or context is no longer valid.
 func (ar *AppendResult) UpdatedSchema(ctx context.Context) (*storagepb.TableSchema, error) {
 	select {
 	case <-ctx.Done():
 		return nil, fmt.Errorf("context done")
 	case <-ar.Ready():
-		return ar.updatedSchema, nil
+		if ar.response != nil {
+			if schema := ar.response.GetUpdatedSchema(); schema != nil {
+				return proto.Clone(schema).(*storagepb.TableSchema), nil
+			}
+		}
+		return nil, nil
 	}
 }
 
@@ -116,9 +185,12 @@ func newPendingWrite(appends [][]byte) *pendingWrite {
 
 // markDone propagates finalization of an append request to the associated
 // AppendResult.
-func (pw *pendingWrite) markDone(startOffset int64, err error, fc *flowController) {
+func (pw *pendingWrite) markDone(resp *storagepb.AppendRowsResponse, err error, fc *flowController) {
+	if resp != nil {
+		pw.result.response = resp
+	}
 	pw.result.err = err
-	pw.result.offset = startOffset
+
 	close(pw.result.ready)
 	// Clear the reference to the request.
 	pw.request = nil

bigquery/storage/managedwriter/appendresult_test.go

@@ -16,9 +16,16 @@ package managedwriter
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"testing"
 	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
+	storagepb "google.golang.org/genproto/googleapis/cloud/bigquery/storage/v1"
+	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 func TestAppendResult(t *testing.T) {
@@ -37,6 +44,7 @@ func TestAppendResult(t *testing.T) {
 }
 
 func TestPendingWrite(t *testing.T) {
+	ctx := context.Background()
 	wantRowData := [][]byte{
 		[]byte("row1"),
 		[]byte("row2"),
@@ -63,9 +71,9 @@ func TestPendingWrite(t *testing.T) {
 	}
 
 	// 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)
+	pending.markDone(&storage.AppendRowsResponse{}, nil, nil)
+	if gotOff := pending.result.offset(ctx); gotOff != NoStreamOffset {
+		t.Errorf("mismatch on completed AppendResult without offset: got %d want %d", gotOff, NoStreamOffset)
 	}
 	if pending.result.err != nil {
 		t.Errorf("mismatch in error on AppendResult, got %v want nil", pending.result.err)
@@ -85,20 +93,30 @@ func TestPendingWrite(t *testing.T) {
 
 	// 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)
+	wantOffset := int64(101)
+	f := WithOffset(wantOffset)
 	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)
+	if pending.request.GetOffset().GetValue() != wantOffset {
+		t.Errorf("offset mismatch, got %d wanted %d", pending.request.GetOffset().GetValue(), wantOffset)
 	}
 
 	// Verify completion behavior with an error.
 	wantErr := fmt.Errorf("foo")
-	pending.markDone(reportedOffset, wantErr, nil)
+
+	testResp := &storagepb.AppendRowsResponse{
+		Response: &storagepb.AppendRowsResponse_AppendResult_{
+			AppendResult: &storagepb.AppendRowsResponse_AppendResult{
+				Offset: &wrapperspb.Int64Value{
+					Value: wantOffset,
+				},
+			},
+		},
+	}
+	pending.markDone(testResp, wantErr, nil)
 
 	if pending.request != nil {
 		t.Errorf("expected request to be cleared, is present: %#v", pending.request)
@@ -118,12 +136,20 @@ func TestPendingWrite(t *testing.T) {
 	case <-time.After(100 * time.Millisecond):
 		t.Errorf("possible blocking on completed AppendResult")
 	case <-pending.result.Ready():
-		if pending.result.offset != reportedOffset {
-			t.Errorf("mismatch on completed AppendResult offset: got %d want %d", pending.result.offset, reportedOffset)
+		gotOffset, gotErr := pending.result.GetResult(ctx)
+		if gotOffset != wantOffset {
+			t.Errorf("GetResult: mismatch on completed AppendResult offset: got %d want %d", gotOffset, wantOffset)
 		}
-		if pending.result.err != wantErr {
-			t.Errorf("mismatch in errors, got %v want %v", pending.result.err, wantErr)
+		if gotErr != wantErr {
+			t.Errorf("GetResult: mismatch in errors, got %v want %v", gotErr, wantErr)
+		}
+		// Now, check FullResponse.
+		gotResp, gotErr := pending.result.FullResponse(ctx)
+		if gotErr != wantErr {
+			t.Errorf("FullResponse: mismatch in errors, got %v want %v", gotErr, wantErr)
+		}
+		if diff := cmp.Diff(gotResp, testResp, protocmp.Transform()); diff != "" {
+			t.Errorf("FullResponse diff: %s", diff)
 		}
 	}
-
 }

bigquery/storage/managedwriter/doc.go

@@ -24,8 +24,7 @@ feature-rich successor to the classic BigQuery streaming interface, which is pre
 in cloud.google.com/go/bigquery, and the tabledata.insertAll method if you're more familiar with the BigQuery v2 REST
 methods.
 
-
-Creating a Client
+# Creating a Client
 
 To start working with this package, create a client:
 
@@ -35,8 +34,7 @@ To start working with this package, create a client:
 		// TODO: Handle error.
 	}
 
-
-Defining the Protocol Buffer Schema
+# Defining the Protocol Buffer Schema
 
 The write functionality of BigQuery Storage requires data to be sent using encoded
 protocol buffer messages using proto2 wire format.  As the protocol buffer is not
@@ -70,7 +68,7 @@ contains functionality to normalize the descriptor into a self-contained definit
 The adapt subpackage also contains functionality for generating a DescriptorProto using
 a BigQuery table's schema directly.
 
-Constructing a ManagedStream
+# Constructing a ManagedStream
 
 The ManagedStream handles management of the underlying write connection to the BigQuery
 Storage service.  You can either create a write session explicitly and pass it in, or
@@ -102,7 +100,7 @@ In addition, NewManagedStream can create new streams implicitly:
 		// TODO: Handle error.
 	}
 
-Writing Data
+# Writing Data
 
 Use the AppendRows function to write one or more serialized proto messages to a stream. You
 can choose to specify an offset in the stream to handle de-duplication for user-created streams,
@@ -111,42 +109,40 @@ but a "default" stream neither accepts nor reports offsets.
 AppendRows returns a future-like object that blocks until the write is successful or yields
 an error.
 
+		// Define a couple of messages.
+		mesgs := []*myprotopackage.MyCompiledMessage{
+			{
+				UserName: proto.String("johndoe"),
+				EmailAddress: proto.String("[email protected]",
+				FavoriteNumbers: []proto.Int64{1,42,12345},
+			},
+			{
+				UserName: proto.String("janesmith"),
+				EmailAddress: proto.String("[email protected]",
+				FavoriteNumbers: []proto.Int64{1,3,5,7,9},
+			},
+		}
 
-	// Define a couple of messages.
-	mesgs := []*myprotopackage.MyCompiledMessage{
-		{
-			UserName: proto.String("johndoe"),
-			EmailAddress: proto.String("[email protected]",
-			FavoriteNumbers: []proto.Int64{1,42,12345},
-		},
-		{
-			UserName: proto.String("janesmith"),
-			EmailAddress: proto.String("[email protected]",
-			FavoriteNumbers: []proto.Int64{1,3,5,7,9},
-		},
-	}
-
-	// Encode the messages into binary format.
-	encoded := make([][]byte, len(mesgs))
-	for k, v := range mesgs{
-		b, err := proto.Marshal(v)
+		// Encode the messages into binary format.
+		encoded := make([][]byte, len(mesgs))
+		for k, v := range mesgs{
+			b, err := proto.Marshal(v)
+			if err != nil {
+				// TODO: Handle error.
+			}
+			encoded[k] = b
+	 	}
+
+		// Send the rows to the service, and specify an offset for managing deduplication.
+		result, err := managedStream.AppendRows(ctx, encoded, WithOffset(0))
+
+		// Block until the write is complete and return the result.
+		returnedOffset, err := result.GetResult(ctx)
 		if err != nil {
 			// TODO: Handle error.
 		}
-		encoded[k] = b
- 	}
 
-	// Send the rows to the service, and specify an offset for managing deduplication.
-	result, err := managedStream.AppendRows(ctx, encoded, WithOffset(0))
-
-	// Block until the write is complete and return the result.
-	returnedOffset, err := result.GetResult(ctx)
-	if err != nil {
-		// TODO: Handle error.
-	}
-
-
-Buffered Stream Management
+# Buffered Stream Management
 
 For Buffered streams, users control when data is made visible in the destination table/stream
 independently of when it is written.  Use FlushRows on the ManagedStream to advance the flush
@@ -156,12 +152,11 @@ point ahead in the stream.
 	// ahead to make the first 1000 rows available.
 	flushOffset, err := managedStream.FlushRows(ctx, 1000)
 
-Pending Stream Management
+# Pending Stream Management
 
 Pending streams allow users to commit data from multiple streams together once the streams
 have been finalized, meaning they'll no longer allow further data writes.
 
-
 	// First, finalize the stream we're writing into.
 	totalRows, err := managedStream.Finalize(ctx)
 	if err != nil {
@@ -175,6 +170,5 @@ have been finalized, meaning they'll no longer allow further data writes.
 	// Using the client, we can commit data from multple streams to the same
 	// table atomically.
 	resp, err := client.BatchCommitWriteStreams(ctx, req)
-
 */
 package managedwriter