BigQuery: Can iterate over QueryJob results without explicitly calling result() (googleapis#4350)

alixhami · web-flow · commit d3999a867f30 · 2017-11-13T14:21:50.000-08:00
diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py
@@ -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.
diff --git a/bigquery/tests/system.py b/bigquery/tests/system.py
@@ -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',
diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py
@@ -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'
