Skip to content

Commit 30de15f

Browse files
authored
perf: use getQueryResults from DB-API (#375)
I suspect `list_rows` (`tabledata.list`) was being called directly due to no `page_size` parameter on `QueryJob.result` at the time.
1 parent 0c387da commit 30de15f

File tree

2 files changed

+7
-8
lines changed

2 files changed

+7
-8
lines changed

google/cloud/bigquery/dbapi/cursor.py

+1-6
Original file line numberDiff line numberDiff line change
@@ -229,19 +229,14 @@ def _try_fetch(self, size=None):
229229
return
230230

231231
if self._query_data is None:
232-
client = self.connection._client
233232
bqstorage_client = self.connection._bqstorage_client
234233

235234
if bqstorage_client is not None:
236235
rows_iterable = self._bqstorage_fetch(bqstorage_client)
237236
self._query_data = _helpers.to_bq_table_rows(rows_iterable)
238237
return
239238

240-
rows_iter = client.list_rows(
241-
self._query_job.destination,
242-
selected_fields=self._query_job._query_results.schema,
243-
page_size=self.arraysize,
244-
)
239+
rows_iter = self._query_job.result(page_size=self.arraysize)
245240
self._query_data = iter(rows_iter)
246241

247242
def _bqstorage_fetch(self, bqstorage_client):

tests/unit/test_dbapi_cursor.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,8 @@ def _mock_client(
6666
num_dml_affected_rows=num_dml_affected_rows,
6767
dry_run=dry_run_job,
6868
total_bytes_processed=total_bytes_processed,
69+
rows=rows,
6970
)
70-
mock_client.list_rows.return_value = rows
7171
mock_client._default_query_job_config = default_query_job_config
7272

7373
# Assure that the REST client gets used, not the BQ Storage client.
@@ -102,9 +102,13 @@ def _mock_job(
102102
num_dml_affected_rows=None,
103103
dry_run=False,
104104
total_bytes_processed=0,
105+
rows=None,
105106
):
106107
from google.cloud.bigquery import job
107108

109+
if rows is None:
110+
rows = []
111+
108112
mock_job = mock.create_autospec(job.QueryJob)
109113
mock_job.error_result = None
110114
mock_job.state = "DONE"
@@ -114,7 +118,7 @@ def _mock_job(
114118
mock_job.result.side_effect = exceptions.NotFound
115119
mock_job.total_bytes_processed = total_bytes_processed
116120
else:
117-
mock_job.result.return_value = mock_job
121+
mock_job.result.return_value = rows
118122
mock_job._query_results = self._mock_results(
119123
total_rows=total_rows,
120124
schema=schema,

0 commit comments

Comments
 (0)