Return int from 'total_rows'/'total_bytes_processed', if present.

Closes googleapis#3004.
richkadel · Feb 13, 2017 · 7f83dc5 · 7f83dc5
1 parent 871a676
commit 7f83dc5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/bigquery/google/cloud/bigquery/query.py b/bigquery/google/cloud/bigquery/query.py
@@ -208,7 +208,9 @@ def total_rows(self):
         :rtype: int, or ``NoneType``
         :returns: Count generated on the server (None until set by the server).
         """
-        return self._properties.get('totalRows')
+        total_rows = self._properties.get('totalRows')
+        if total_rows is not None:
+            return int(total_rows)
 
     @property
     def total_bytes_processed(self):
@@ -220,7 +222,9 @@ def total_bytes_processed(self):
         :rtype: int, or ``NoneType``
         :returns: Count generated on the server (None until set by the server).
         """
-        return self._properties.get('totalBytesProcessed')
+        total_bytes_processed = self._properties.get('totalBytesProcessed')
+        if total_bytes_processed is not None:
+            return int(total_bytes_processed)
 
     @property
     def rows(self):

diff --git a/bigquery/unit_tests/test_query.py b/bigquery/unit_tests/test_query.py
@@ -124,9 +124,15 @@ def _verifyResourceProperties(self, query, resource):
         self.assertEqual(query.complete, resource.get('jobComplete'))
         self.assertEqual(query.errors, resource.get('errors'))
         self.assertEqual(query.page_token, resource.get('pageToken'))
-        self.assertEqual(query.total_rows, resource.get('totalRows'))
-        self.assertEqual(query.total_bytes_processed,
-                         resource.get('totalBytesProcessed'))
+        if 'totalRows' in resource:
+            self.assertEqual(query.total_rows, int(resource['totalRows']))
+        else:
+            self.assertIsNone(query.total_rows)
+        if 'totalBytesProcessed' in resource:
+            self.assertEqual(query.total_bytes_processed,
+                             int(resource['totalBytesProcessed']))
+        else:
+            self.assertIsNone(query.total_bytes_processed)
 
         if 'jobReference' in resource:
             self.assertEqual(query.name, resource['jobReference']['jobId'])
@@ -336,6 +342,14 @@ def test_total_rows_present_integer(self):
         query._set_properties(resource)
         self.assertEqual(query.total_rows, TOTAL_ROWS)
 
+    def test_total_rows_present_string(self):
+        TOTAL_ROWS = 42
+        client = _Client(self.PROJECT)
+        query = self._make_one(self.QUERY, client)
+        resource = {'totalRows': str(TOTAL_ROWS)}
+        query._set_properties(resource)
+        self.assertEqual(query.total_rows, TOTAL_ROWS)
+
     def test_total_bytes_processed_missing(self):
         client = _Client(self.PROJECT)
         query = self._make_one(self.QUERY, client)
@@ -349,6 +363,14 @@ def test_total_bytes_processed_present_integer(self):
         query._set_properties(resource)
         self.assertEqual(query.total_bytes_processed, TOTAL_BYTES_PROCESSED)
 
+    def test_total_bytes_processed_present_string(self):
+        TOTAL_BYTES_PROCESSED = 123456
+        client = _Client(self.PROJECT)
+        query = self._make_one(self.QUERY, client)
+        resource = {'totalBytesProcessed': str(TOTAL_BYTES_PROCESSED)}
+        query._set_properties(resource)
+        self.assertEqual(query.total_bytes_processed, TOTAL_BYTES_PROCESSED)
+
     def test_schema(self):
         client = _Client(self.PROJECT)
         query = self._make_one(self.QUERY, client)