Add 'Query.num_dml_affected_rows' property. (#3460)

Read-only, set from servier-provided 'numDmlAffectedRows' field.

Closes #2920.

Author: tseaver

bigquery/google/cloud/bigquery/query.py

@@ -226,6 +226,20 @@ def total_bytes_processed(self):
         if total_bytes_processed is not None:
             return int(total_bytes_processed)
 
+    @property
+    def num_dml_affected_rows(self):
+        """Total number of rows affected by a DML query.
+
+        See:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#numDmlAffectedRows
+
+        :rtype: int, or ``NoneType``
+        :returns: Count generated on the server (None until set by the server).
+        """
+        num_dml_affected_rows = self._properties.get('numDmlAffectedRows')
+        if num_dml_affected_rows is not None:
+            return int(num_dml_affected_rows)
+
     @property
     def rows(self):
         """Query results.

bigquery/tests/unit/test_query.py

@@ -70,6 +70,7 @@ def _makeResource(self, complete=False):
             ]
             resource['pageToken'] = self.TOKEN
             resource['totalBytesProcessed'] = 100000
+            resource['numDmlAffectedRows'] = 123
             resource['cacheHit'] = False
 
         return resource
@@ -124,10 +125,12 @@ 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'))
+
         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']))
@@ -139,6 +142,12 @@ def _verifyResourceProperties(self, query, resource):
         else:
             self.assertIsNone(query.name)
 
+        if 'numDmlAffectedRows' in resource:
+            self.assertEqual(query.num_dml_affected_rows,
+                             int(resource['numDmlAffectedRows']))
+        else:
+            self.assertIsNone(query.num_dml_affected_rows)
+
         self._verify_udf_resources(query, resource)
         self._verifyQueryParameters(query, resource)
         self._verifySchema(query, resource)
@@ -371,6 +380,27 @@ def test_total_bytes_processed_present_string(self):
         query._set_properties(resource)
         self.assertEqual(query.total_bytes_processed, TOTAL_BYTES_PROCESSED)
 
+    def test_num_dml_affected_rows_missing(self):
+        client = _Client(self.PROJECT)
+        query = self._make_one(self.QUERY, client)
+        self.assertIsNone(query.num_dml_affected_rows)
+
+    def test_num_dml_affected_rows_present_integer(self):
+        DML_AFFECTED_ROWS = 123456
+        client = _Client(self.PROJECT)
+        query = self._make_one(self.QUERY, client)
+        resource = {'numDmlAffectedRows': DML_AFFECTED_ROWS}
+        query._set_properties(resource)
+        self.assertEqual(query.num_dml_affected_rows, DML_AFFECTED_ROWS)
+
+    def test_num_dml_affected_rows_present_string(self):
+        DML_AFFECTED_ROWS = 123456
+        client = _Client(self.PROJECT)
+        query = self._make_one(self.QUERY, client)
+        resource = {'numDmlAffectedRows': str(DML_AFFECTED_ROWS)}
+        query._set_properties(resource)
+        self.assertEqual(query.num_dml_affected_rows, DML_AFFECTED_ROWS)
+
     def test_schema(self):
         client = _Client(self.PROJECT)
         query = self._make_one(self.QUERY, client)