Ensure total row count is integer or None for 'Query.fetch_data'.

tseaver · tseaver · commit 7d829ba94fbe · 2016-07-14T10:25:41.000-04:00
diff --git a/gcloud/bigquery/query.py b/gcloud/bigquery/query.py
@@ -343,6 +343,8 @@ def fetch_data(self, max_results=None, page_token=None, start_index=None,
         self._set_properties(response)
 
         total_rows = response.get('totalRows')
+        if total_rows is not None:
+            total_rows = int(total_rows)
         page_token = response.get('pageToken')
         rows_data = _rows_from_json(response.get('rows', ()), self.schema)
 
diff --git a/gcloud/bigquery/test_query.py b/gcloud/bigquery/test_query.py
@@ -47,7 +47,7 @@ def _makeResource(self, complete=False):
         }
 
         if complete:
-            resource['totalRows'] = 1000
+            resource['totalRows'] = '1000'
             resource['rows'] = [
                 {'f': [
                     {'v': 'Phred Phlyntstone'},
@@ -240,6 +240,7 @@ def test_fetch_data_w_bound_client(self):
         PATH = 'projects/%s/queries/%s' % (self.PROJECT, self.JOB_NAME)
         BEFORE = self._makeResource(complete=False)
         AFTER = self._makeResource(complete=True)
+        del AFTER['totalRows']
 
         conn = _Connection(AFTER)
         client = _Client(project=self.PROJECT, connection=conn)
@@ -255,7 +256,7 @@ def test_fetch_data_w_bound_client(self):
         self.assertEqual(rows[1], ('Bharney Rhubble', 33))
         self.assertEqual(rows[2], ('Wylma Phlyntstone', 29))
         self.assertEqual(rows[3], ('Bhettye Rhubble', 27))
-        self.assertEqual(total_rows, AFTER['totalRows'])
+        self.assertEqual(total_rows, None)
         self.assertEqual(page_token, AFTER['pageToken'])
 
         self.assertEqual(len(conn._requested), 1)
@@ -290,7 +291,7 @@ def test_fetch_data_w_alternate_client(self):
         self.assertEqual(rows[1], ('Bharney Rhubble', 33))
         self.assertEqual(rows[2], ('Wylma Phlyntstone', 29))
         self.assertEqual(rows[3], ('Bhettye Rhubble', 27))
-        self.assertEqual(total_rows, AFTER['totalRows'])
+        self.assertEqual(total_rows, int(AFTER['totalRows']))
         self.assertEqual(page_token, AFTER['pageToken'])
 
         self.assertEqual(len(conn1._requested), 0)
