Skip to content

Commit

Permalink
feat(bigquery): expose start index parameter for query result (#121)
Browse files Browse the repository at this point in the history
* feat(bigquery): expose start index parameter for query result

* feat(bigquery): nit
  • Loading branch information
HemangChothani authored Jun 6, 2020
1 parent cfaa901 commit be86de3
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 1 deletion.
10 changes: 9 additions & 1 deletion google/cloud/bigquery/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -3160,7 +3160,12 @@ def _begin(self, client=None, retry=DEFAULT_RETRY, timeout=None):
raise

def result(
self, page_size=None, max_results=None, retry=DEFAULT_RETRY, timeout=None
self,
page_size=None,
max_results=None,
retry=DEFAULT_RETRY,
timeout=None,
start_index=None,
):
"""Start the job and wait for it to complete and get the result.
Expand All @@ -3177,6 +3182,8 @@ def result(
before using ``retry``.
If multiple requests are made under the hood, ``timeout``
applies to each individual request.
start_index (Optional[int]):
The zero-based index of the starting row to read.
Returns:
google.cloud.bigquery.table.RowIterator:
Expand Down Expand Up @@ -3230,6 +3237,7 @@ def result(
dest_table,
page_size=page_size,
max_results=max_results,
start_index=start_index,
retry=retry,
timeout=timeout,
)
Expand Down
12 changes: 12 additions & 0 deletions tests/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -1577,6 +1577,18 @@ def test_query_w_page_size(self):
iterator = query_job.result(page_size=page_size)
self.assertEqual(next(iterator.pages).num_items, page_size)

def test_query_w_start_index(self):
start_index = 164652
query_job = Config.CLIENT.query(
"SELECT word FROM `bigquery-public-data.samples.shakespeare`;",
job_id_prefix="test_query_w_start_index_",
)
result1 = query_job.result(start_index=start_index)
total_rows = result1.total_rows

self.assertEqual(result1.extra_params["startIndex"], start_index)
self.assertEqual(len(list(result1)), total_rows - start_index)

def test_query_statistics(self):
"""
A system test to exercise some of the extended query statistics.
Expand Down
40 changes: 40 additions & 0 deletions tests/unit/test_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -4666,6 +4666,46 @@ def test_result_w_page_size(self):
]
)

def test_result_with_start_index(self):
from google.cloud.bigquery.table import RowIterator

query_resource = {
"jobComplete": True,
"jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
"schema": {"fields": [{"name": "col1", "type": "STRING"}]},
"totalRows": "5",
}
tabledata_resource = {
"totalRows": "5",
"pageToken": None,
"rows": [
{"f": [{"v": "abc"}]},
{"f": [{"v": "def"}]},
{"f": [{"v": "ghi"}]},
{"f": [{"v": "jkl"}]},
],
}
connection = _make_connection(query_resource, tabledata_resource)
client = _make_client(self.PROJECT, connection=connection)
resource = self._make_resource(ended=True)
job = self._get_target_class().from_api_repr(resource, client)

start_index = 1

result = job.result(start_index=start_index)

self.assertIsInstance(result, RowIterator)
self.assertEqual(result.total_rows, 5)

rows = list(result)

self.assertEqual(len(rows), 4)
self.assertEqual(len(connection.api_request.call_args_list), 2)
tabledata_list_request = connection.api_request.call_args_list[1]
self.assertEqual(
tabledata_list_request[1]["query_params"]["startIndex"], start_index
)

def test_result_error(self):
from google.cloud import exceptions

Expand Down

0 comments on commit be86de3

Please sign in to comment.