Merge be6c5b4 into 3a0881d

src/commons/big_query/big_query.py

@@ -111,15 +111,15 @@ def execute_query(self, query, use_legacy_sql=True):
         page_token = None
 
         while True:
-            page = self.service.jobs().getQueryResults(
+            response = self.service.jobs().getQueryResults(
                 pageToken=page_token,
                 **query_job['jobReference']).execute(num_retries=2)
-
-            results.extend(page.get('rows', []))
-
-            page_token = page.get('pageToken')
-            if not page_token:
-                break
+            job_complete = response.get('jobComplete')
+            if job_complete:
+                results.extend(response.get('rows', []))
+                page_token = response.get('pageToken')
+                if not page_token:
+                    break
 
         return results
 

tests/commons/big_query/test_big_query.py

@@ -184,6 +184,46 @@ def __create_tables_list_responses():
              content('tests/json_samples/bigquery_table_list_page_last.json'))
         ])
 
+    def test_execute_query_when_executing_long_query(self):
+        # given
+        self._create_http.return_value = self.__execute_long_query_responses()
+        # when
+        result = BigQuery().execute_query("SELECT * FROM tableXYZ")
+        # then
+        self.assertEqual(result, [
+    {
+      "f": [
+        {
+          "v": "local-project-bbq"
+        },
+        {
+          "v": "test1"
+        }
+      ]
+    },
+    {
+      "f": [
+        {
+          "v": "project-bbq-restoration"
+        },
+        {
+          "v": "smoke_test_US"
+        }
+      ]
+    }
+        ])
+
+
+    @staticmethod
+    def __execute_long_query_responses():
+        return HttpMockSequence([
+            ({'status': '200'}, content('tests/json_samples/big_query/discovery_v1_apis_bigquery_v2_rest.json')),
+            ({'status': '200'}, content('tests/json_samples/big_query/query_response.json')),
+            ({'status': '200'}, content('tests/json_samples/big_query/get_query_results_job_not_completed.json')),
+            ({'status': '200'}, content('tests/json_samples/big_query/get_query_results_job_completed.json'))
+        ])
+
+
     @staticmethod
     def __create_tables_list_responses_with_503():
         return HttpMockSequence([