feat(bigquery/storage/managedwriter): add state tracking

bigquery/storage/managedwriter/appendresult.go

@@ -22,6 +22,10 @@ import (
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
+// NoOffset is a sentinel value for signalling we're not tracking
+// stream offset (e.g. a default stream which allows simultaneous append streams).
+const NoOffset int64 = -1
+
 // AppendResult tracks the status of a single row of data.
 type AppendResult struct {
 	// rowData contains the serialized row data.
@@ -49,12 +53,6 @@ func (ar *AppendResult) Ready() <-chan struct{} { return ar.ready }
 // GetResult returns the optional offset of this row, or the associated
 // error.
 func (ar *AppendResult) GetResult(ctx context.Context) (int64, error) {
-	select {
-	case <-ar.Ready():
-		return ar.offset, ar.err
-	default:
-	}
-
 	select {
 	case <-ctx.Done():
 		return 0, ctx.Err()
@@ -106,10 +104,6 @@ func newPendingWrite(appends [][]byte, offset int64) *pendingWrite {
 	return pw
 }
 
-// noOffsetSpecified is a sentinel value for signalling we're not tracking
-// stream offset (e.g. a default stream which allows simultaneous append streams).
-var noOffsetSpecified int64 = -1
-
 // markDone propagates finalization of an append request to associated
 // AppendResult references.
 func (pw *pendingWrite) markDone(startOffset int64, err error) {

bigquery/storage/managedwriter/appendresult_test.go

@@ -41,10 +41,19 @@ func TestPendingWrite(t *testing.T) {
 	var wantOffset int64 = 99
 
 	// first, verify no offset behavior
-	pending := newPendingWrite(wantRowData, noOffsetSpecified)
+	pending := newPendingWrite(wantRowData, NoOffset)
 	if pending.request.GetOffset() != nil {
 		t.Errorf("request should have no offset, but is present: %q", pending.request.GetOffset().GetValue())
 	}
+	pending.markDone(NoOffset, nil)
+	for k, ar := range pending.results {
+		if ar.offset != NoOffset {
+			t.Errorf("mismatch on completed AppendResult(%d) without offset: got %d want %d", k, ar.offset, NoOffset)
+		}
+		if ar.err != nil {
+			t.Errorf("mismatch in error on AppendResult(%d), got %v want nil", k, ar.err)
+		}
+	}
 
 	// now, verify behavior with a valid offset
 	pending = newPendingWrite(wantRowData, 99)