Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

perf: use getQueryResults from DB-API #375

Merged
merged 1 commit into from Nov 10, 2020
Merged

perf: use getQueryResults from DB-API #375

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
7 changes: 1 addition & 6 deletions google/cloud/bigquery/dbapi/cursor.py
View file
Open in desktop
Expand Up @@ -229,19 +229,14 @@ def _try_fetch(self, size=None):
return

if self._query_data is None:
client = self.connection._client
bqstorage_client = self.connection._bqstorage_client

if bqstorage_client is not None:
rows_iterable = self._bqstorage_fetch(bqstorage_client)
self._query_data = _helpers.to_bq_table_rows(rows_iterable)
return

rows_iter = client.list_rows(
self._query_job.destination,
selected_fields=self._query_job._query_results.schema,
page_size=self.arraysize,
)
rows_iter = self._query_job.result(page_size=self.arraysize)
self._query_data = iter(rows_iter)

def _bqstorage_fetch(self, bqstorage_client):
Expand Down
8 changes: 6 additions & 2 deletions tests/unit/test_dbapi_cursor.py
View file
Open in desktop
Expand Up @@ -66,8 +66,8 @@ def _mock_client(
num_dml_affected_rows=num_dml_affected_rows,
dry_run=dry_run_job,
total_bytes_processed=total_bytes_processed,
rows=rows,
)
mock_client.list_rows.return_value = rows
mock_client._default_query_job_config = default_query_job_config

# Assure that the REST client gets used, not the BQ Storage client.
Expand Down Expand Up @@ -102,9 +102,13 @@ def _mock_job(
num_dml_affected_rows=None,
dry_run=False,
total_bytes_processed=0,
rows=None,
):
from google.cloud.bigquery import job

if rows is None:
rows = []

mock_job = mock.create_autospec(job.QueryJob)
mock_job.error_result = None
mock_job.state = "DONE"
Expand All @@ -114,7 +118,7 @@ def _mock_job(
mock_job.result.side_effect = exceptions.NotFound
mock_job.total_bytes_processed = total_bytes_processed
else:
mock_job.result.return_value = mock_job
mock_job.result.return_value = rows
mock_job._query_results = self._mock_results(
total_rows=total_rows,
schema=schema,
Expand Down