File tree 2 files changed +7
-8
lines changed
google/cloud/bigquery/dbapi
2 files changed +7
-8
lines changed Original file line number Diff line number Diff line change @@ -229,19 +229,14 @@ def _try_fetch(self, size=None):
229
229
return
230
230
231
231
if self ._query_data is None :
232
- client = self .connection ._client
233
232
bqstorage_client = self .connection ._bqstorage_client
234
233
235
234
if bqstorage_client is not None :
236
235
rows_iterable = self ._bqstorage_fetch (bqstorage_client )
237
236
self ._query_data = _helpers .to_bq_table_rows (rows_iterable )
238
237
return
239
238
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 )
245
240
self ._query_data = iter (rows_iter )
246
241
247
242
def _bqstorage_fetch (self , bqstorage_client ):
Original file line number Diff line number Diff line change @@ -66,8 +66,8 @@ def _mock_client(
66
66
num_dml_affected_rows = num_dml_affected_rows ,
67
67
dry_run = dry_run_job ,
68
68
total_bytes_processed = total_bytes_processed ,
69
+ rows = rows ,
69
70
)
70
- mock_client .list_rows .return_value = rows
71
71
mock_client ._default_query_job_config = default_query_job_config
72
72
73
73
# Assure that the REST client gets used, not the BQ Storage client.
@@ -102,9 +102,13 @@ def _mock_job(
102
102
num_dml_affected_rows = None ,
103
103
dry_run = False ,
104
104
total_bytes_processed = 0 ,
105
+ rows = None ,
105
106
):
106
107
from google .cloud .bigquery import job
107
108
109
+ if rows is None :
110
+ rows = []
111
+
108
112
mock_job = mock .create_autospec (job .QueryJob )
109
113
mock_job .error_result = None
110
114
mock_job .state = "DONE"
@@ -114,7 +118,7 @@ def _mock_job(
114
118
mock_job .result .side_effect = exceptions .NotFound
115
119
mock_job .total_bytes_processed = total_bytes_processed
116
120
else :
117
- mock_job .result .return_value = mock_job
121
+ mock_job .result .return_value = rows
118
122
mock_job ._query_results = self ._mock_results (
119
123
total_rows = total_rows ,
120
124
schema = schema ,
You can’t perform that action at this time.
0 commit comments