fix(bigquery): fix start index with page size for list rows (#27)

Co-authored-by: Peter Lamut <plamut@users.noreply.github.com>
googleapis · Mar 4, 2020 · 400673b · 400673b
1 parent c88a6dd
commit 400673b
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 0 deletions.
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
@@ -1367,6 +1367,8 @@ def _get_next_page_response(self):
         """
         params = self._get_query_params()
         if self._page_size is not None:
+            if self.page_number and "startIndex" in params:
+                del params["startIndex"]
             params["maxResults"] = self._page_size
         return self.api_request(
             method=self._HTTP_METHOD, path=self.path, query_params=params

diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
@@ -5612,6 +5612,71 @@ def _bigquery_timestamp_float_repr(ts_float):
             method="GET", path="/%s" % PATH, query_params={}, timeout=7.5
         )
 
+    def test_list_rows_w_start_index_w_page_size(self):
+        from google.cloud.bigquery.schema import SchemaField
+        from google.cloud.bigquery.table import Table
+        from google.cloud.bigquery.table import Row
+
+        PATH = "projects/%s/datasets/%s/tables/%s/data" % (
+            self.PROJECT,
+            self.DS_ID,
+            self.TABLE_ID,
+        )
+
+        page_1 = {
+            "totalRows": 4,
+            "pageToken": "some-page-token",
+            "rows": [
+                {"f": [{"v": "Phred Phlyntstone"}]},
+                {"f": [{"v": "Bharney Rhubble"}]},
+            ],
+        }
+        page_2 = {
+            "totalRows": 4,
+            "rows": [
+                {"f": [{"v": "Wylma Phlyntstone"}]},
+                {"f": [{"v": "Bhettye Rhubble"}]},
+            ],
+        }
+        creds = _make_credentials()
+        http = object()
+        client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+        conn = client._connection = make_connection(page_1, page_2)
+        full_name = SchemaField("full_name", "STRING", mode="REQUIRED")
+        table = Table(self.TABLE_REF, schema=[full_name])
+        iterator = client.list_rows(table, max_results=4, page_size=2, start_index=1)
+        pages = iterator.pages
+        rows = list(six.next(pages))
+        extra_params = iterator.extra_params
+        f2i = {"full_name": 0}
+        self.assertEqual(len(rows), 2)
+        self.assertEqual(rows[0], Row(("Phred Phlyntstone",), f2i))
+        self.assertEqual(rows[1], Row(("Bharney Rhubble",), f2i))
+
+        rows = list(six.next(pages))
+
+        self.assertEqual(len(rows), 2)
+        self.assertEqual(rows[0], Row(("Wylma Phlyntstone",), f2i))
+        self.assertEqual(rows[1], Row(("Bhettye Rhubble",), f2i))
+        self.assertEqual(extra_params, {"startIndex": 1})
+
+        conn.api_request.assert_has_calls(
+            [
+                mock.call(
+                    method="GET",
+                    path="/%s" % PATH,
+                    query_params={"startIndex": 1, "maxResults": 2},
+                    timeout=None,
+                ),
+                mock.call(
+                    method="GET",
+                    path="/%s" % PATH,
+                    query_params={"pageToken": "some-page-token", "maxResults": 2},
+                    timeout=None,
+                ),
+            ]
+        )
+
     def test_list_rows_empty_table(self):
         response = {"totalRows": "0", "rows": []}
         creds = _make_credentials()