diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go index 24a3e20fa8d3..ddb8f14b0688 100644 --- a/bigtable/bigtable.go +++ b/bigtable/bigtable.go @@ -161,7 +161,10 @@ func (c *Client) Open(table string) *Table { return &Table{ c: c, table: table, - md: metadata.Pairs(resourcePrefixHeader, c.fullTableName(table), requestParamsHeader, c.requestParamsHeaderValue(table)), + md: metadata.Join(metadata.Pairs( + resourcePrefixHeader, c.fullTableName(table), + requestParamsHeader, c.requestParamsHeaderValue(table), + ), btopt.WithFeatureFlags()), } } @@ -246,6 +249,10 @@ func (t *Table) ReadRows(ctx context.Context, arg RowSet, f func(Row) bool, opts } } + if res.LastScannedRowKey != nil { + prevRowKey = string(res.LastScannedRowKey) + } + // Handle any incoming RequestStats. This should happen at most once. if res.RequestStats != nil && settings.fullReadStatsFunc != nil { stats := makeFullReadStats(res.RequestStats) diff --git a/bigtable/retry_test.go b/bigtable/retry_test.go index 3df436e93ecc..54565e97fefb 100644 --- a/bigtable/retry_test.go +++ b/bigtable/retry_test.go @@ -413,6 +413,15 @@ func TestRetryReadRows(t *testing.T) { case 3: // Write two more rows must(writeReadRowsResponse(ss, "c", "d")) + err = status.Errorf(codes.Unavailable, "") + case 4: + must(ss.SendMsg(&btpb.ReadRowsResponse{LastScannedRowKey: []byte("e")})) + err = status.Errorf(codes.Unavailable, "") + case 5: + if want, got := "e\x00", string(req.Rows.RowRanges[0].GetStartKeyClosed()); want != got { + t.Errorf("3 range retries: got %q, want %q", got, want) + } + must(writeReadRowsResponse(ss, "f", "g")) err = nil } errCount++ @@ -424,7 +433,7 @@ func TestRetryReadRows(t *testing.T) { got = append(got, r.Key()) return true })) - want := []string{"a", "b", "c", "d"} + want := []string{"a", "b", "c", "d", "f", "g"} if !testutil.Equal(got, want) { t.Errorf("retry range integration: got %v, want %v", got, want) }