Skip to content

Commit

Permalink
BigQuery: Can iterate over QueryJob results without explicitly callin…
Browse files Browse the repository at this point in the history
…g result() (#4350)
  • Loading branch information
alixhami committed Nov 13, 2017
1 parent a105bba commit d3999a8
Show file tree
Hide file tree
Showing 3 changed files with 29 additions and 0 deletions.
3 changes: 3 additions & 0 deletions bigquery/google/cloud/bigquery/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -1949,6 +1949,9 @@ def result(self, timeout=None, retry=DEFAULT_RETRY):
return self._client.list_rows(dest_table, selected_fields=schema,
retry=retry)

def __iter__(self):
return iter(self.result())


class QueryPlanEntryStep(object):
"""Map a single step in a query plan entry.
Expand Down
7 changes: 7 additions & 0 deletions bigquery/tests/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -1235,6 +1235,13 @@ def test_query_future(self):
row_tuples = [r.values() for r in iterator]
self.assertEqual(row_tuples, [(1,)])

def test_query_iter(self):
import types
query_job = Config.CLIENT.query('SELECT 1')
self.assertIsInstance(iter(query_job), types.GeneratorType)
row_tuples = [r.values() for r in query_job]
self.assertEqual(row_tuples, [(1,)])

def test_query_table_def(self):
gs_url = self._write_csv_to_storage(
'bq_external_test' + unique_resource_id(), 'person_ages.csv',
Expand Down
19 changes: 19 additions & 0 deletions bigquery/tests/unit/test_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -2720,6 +2720,25 @@ def test_reload_w_alternate_client(self):
self.assertEqual(req['path'], PATH)
self._verifyResourceProperties(job, RESOURCE)

def test_iter(self):
import types

begun_resource = self._makeResource()
query_resource = {
'jobComplete': True,
'jobReference': {
'projectId': self.PROJECT,
'jobId': self.JOB_ID,
},
}
done_resource = copy.deepcopy(begun_resource)
done_resource['status'] = {'state': 'DONE'}
connection = _Connection(begun_resource, query_resource, done_resource)
client = _make_client(project=self.PROJECT, connection=connection)
job = self._make_one(self.JOB_ID, self.QUERY, client)

self.assertIsInstance(iter(job), types.GeneratorType)


class TestQueryPlanEntryStep(unittest.TestCase, _Base):
KIND = 'KIND'
Expand Down

0 comments on commit d3999a8

Please sign in to comment.