From cd36506d377d2d5199402a58360c23ba4ce9a3d4 Mon Sep 17 00:00:00 2001
From: Baha Aiman <bahaaiman@google.com>
Date: Wed, 13 Mar 2024 18:19:59 +0530
Subject: [PATCH] fix(bigtable): Accept nil RowSet to read all rows (#9327)

---
 bigtable/bigtable.go                 | 21 ++++++++++++++-------
 bigtable/internal/testproxy/proxy.go |  7 -------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/bigtable/bigtable.go b/bigtable/bigtable.go
index fc39e8f43c3a..0313070c8cfb 100644
--- a/bigtable/bigtable.go
+++ b/bigtable/bigtable.go
@@ -184,16 +184,19 @@ func (t *Table) ReadRows(ctx context.Context, arg RowSet, f func(Row) bool, opts
 	var prevRowKey string
 	attrMap := make(map[string]interface{})
 	err = gax.Invoke(ctx, func(ctx context.Context, _ gax.CallSettings) error {
-		if !arg.valid() {
-			// Empty row set, no need to make an API call.
-			// NOTE: we must return early if arg == RowList{} because reading
-			// an empty RowList from bigtable returns all rows from that table.
-			return nil
-		}
 		req := &btpb.ReadRowsRequest{
 			TableName:    t.c.fullTableName(t.table),
 			AppProfileId: t.c.appProfile,
-			Rows:         arg.proto(),
+		}
+
+		if arg != nil {
+			if !arg.valid() {
+				// Empty row set, no need to make an API call.
+				// NOTE: we must return early if arg == RowList{} because reading
+				// an empty RowList from bigtable returns all rows from that table.
+				return nil
+			}
+			req.Rows = arg.proto()
 		}
 		settings := makeReadSettings(req)
 		for _, opt := range opts {
@@ -222,6 +225,10 @@ func (t *Table) ReadRows(ctx context.Context, arg RowSet, f func(Row) bool, opts
 			}
 			if err != nil {
 				// Reset arg for next Invoke call.
+				if arg == nil {
+					// Should be lowest possible key value, an empty byte array
+					arg = InfiniteRange("")
+				}
 				if req.Reversed {
 					arg = arg.retainRowsBefore(prevRowKey)
 				} else {
diff --git a/bigtable/internal/testproxy/proxy.go b/bigtable/internal/testproxy/proxy.go
index 4738a9cecd11..d5bf11be4753 100644
--- a/bigtable/internal/testproxy/proxy.go
+++ b/bigtable/internal/testproxy/proxy.go
@@ -566,13 +566,6 @@ func (s *goTestProxyServer) ReadRows(ctx context.Context, req *pb.ReadRowsReques
 	rowPbs := rrq.Rows
 	rs := rowSetFromProto(rowPbs)
 
-	// Bigtable client doesn't have a Table.GetAll() function--RowSet must be
-	// provided for ReadRows. Use InfiniteRange() to get the full table.
-	if rs == nil {
-		// Should be lowest possible key value, an empty byte array
-		rs = bigtable.InfiniteRange("")
-	}
-
 	ctx, cancel := btc.timeout(ctx)
 	defer cancel()