feat(bigtable): Support last_scanned_row in the bigtable client (#8345)

steveniemitz · web-flow · commit 961dd38f9e46 · 2023-07-28T16:50:13.000Z
This adds support for last_scanned_row_key to the bigtable client, as well as sending feature flags to the server. note: this depends on #8343
diff --git 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
@@ -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)
 	}

Original file line number	Diff line number	Diff line change
`@@ -161,7 +161,10 @@ func (c Client) Open(table string) Table {`
`161`	`161`	`return &Table{`
`162`	`162`	`c: c,`
`163`	`163`	`table: table,`
`164`		`- md: metadata.Pairs(resourcePrefixHeader, c.fullTableName(table), requestParamsHeader, c.requestParamsHeaderValue(table)),`
	`164`	`+ md: metadata.Join(metadata.Pairs(`
	`165`	`+ resourcePrefixHeader, c.fullTableName(table),`
	`166`	`+ requestParamsHeader, c.requestParamsHeaderValue(table),`
	`167`	`+ ), btopt.WithFeatureFlags()),`
`165`	`168`	`}`
`166`	`169`	`}`
`167`	`170`
`@@ -246,6 +249,10 @@ func (t *Table) ReadRows(ctx context.Context, arg RowSet, f func(Row) bool, opts`
`246`	`249`	`}`
`247`	`250`	`}`
`248`	`251`
	`252`	`+ if res.LastScannedRowKey != nil {`
	`253`	`+ prevRowKey = string(res.LastScannedRowKey)`
	`254`	`+ }`
	`255`	`+`
`249`	`256`	`// Handle any incoming RequestStats. This should happen at most once.`
`250`	`257`	`if res.RequestStats != nil && settings.fullReadStatsFunc != nil {`
`251`	`258`	`stats := makeFullReadStats(res.RequestStats)`