From 9e44ba0815af7991528301f63238dbc540191333 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Wed, 31 May 2017 13:54:09 -0400 Subject: [PATCH] Add 'Query.num_dml_affected_rows' property. (#3460) Read-only, set from servier-provided 'numDmlAffectedRows' field. Closes #2920. --- bigquery/google/cloud/bigquery/query.py | 14 ++++++++++++ bigquery/tests/unit/test_query.py | 30 +++++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/bigquery/google/cloud/bigquery/query.py b/bigquery/google/cloud/bigquery/query.py index 6b764f3c664d5..ee24d8397b736 100644 --- a/bigquery/google/cloud/bigquery/query.py +++ b/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. diff --git a/bigquery/tests/unit/test_query.py b/bigquery/tests/unit/test_query.py index a15833af347d5..c2b3ce5496e18 100644 --- a/bigquery/tests/unit/test_query.py +++ b/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)