feat: add support for more detailed DML stats (#758)

* feat: add support for more detailed DML stats * Move is None check of DmlStats one level higher

Author: plamut

docs/reference.rst

@@ -58,6 +58,7 @@ Job-Related Types
  job.Compression
  job.CreateDisposition
  job.DestinationFormat
+ job.DmlStats
  job.Encoding
  job.OperationType
  job.QueryPlanEntry

google/cloud/bigquery/__init__.py

@@ -56,6 +56,7 @@
 from google.cloud.bigquery.job import CopyJobConfig
 from google.cloud.bigquery.job import CreateDisposition
 from google.cloud.bigquery.job import DestinationFormat
+from google.cloud.bigquery.job import DmlStats
 from google.cloud.bigquery.job import Encoding
 from google.cloud.bigquery.job import ExtractJob
 from google.cloud.bigquery.job import ExtractJobConfig
@@ -142,6 +143,7 @@
  "BigtableOptions",
  "BigtableColumnFamily",
  "BigtableColumn",
+ "DmlStats",
  "CSVOptions",
  "GoogleSheetsOptions",
  "ParquetOptions",

google/cloud/bigquery/job/__init__.py

@@ -31,6 +31,7 @@
 from google.cloud.bigquery.job.load import LoadJob
 from google.cloud.bigquery.job.load import LoadJobConfig
 from google.cloud.bigquery.job.query import _contains_order_by
+from google.cloud.bigquery.job.query import DmlStats
 from google.cloud.bigquery.job.query import QueryJob
 from google.cloud.bigquery.job.query import QueryJobConfig
 from google.cloud.bigquery.job.query import QueryPlanEntry
@@ -66,6 +67,7 @@
  "LoadJob",
  "LoadJobConfig",
  "_contains_order_by",
+ "DmlStats",
  "QueryJob",
  "QueryJobConfig",
  "QueryPlanEntry",

google/cloud/bigquery/job/query.py

@@ -114,6 +114,35 @@ def _to_api_repr_table_defs(value):
  return {k: ExternalConfig.to_api_repr(v) for k, v in value.items()}
 
 
+class DmlStats(typing.NamedTuple):
+ """Detailed statistics for DML statements.
+
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/DmlStats
+ """
+
+ inserted_row_count: int = 0
+ """Number of inserted rows. Populated by DML INSERT and MERGE statements."""
+
+ deleted_row_count: int = 0
+ """Number of deleted rows. populated by DML DELETE, MERGE and TRUNCATE statements.
+ """
+
+ updated_row_count: int = 0
+ """Number of updated rows. Populated by DML UPDATE and MERGE statements."""
+
+ @classmethod
+ def from_api_repr(cls, stats: Dict[str, str]) -> "DmlStats":
+ # NOTE: The field order here must match the order of fields set at the
+ # class level.
+ api_fields = ("insertedRowCount", "deletedRowCount", "updatedRowCount")
+
+ args = (
+ int(stats.get(api_field, default_val))
+ for api_field, default_val in zip(api_fields, cls.__new__.__defaults__)
+ )
+ return cls(*args)
+
+
 class ScriptOptions:
  """Options controlling the execution of scripts.
 
@@ -1079,6 +1108,14 @@ def estimated_bytes_processed(self):
  result = int(result)
  return result
 
+ @property
+ def dml_stats(self) -> Optional[DmlStats]:
+ stats = self._job_statistics().get("dmlStats")
+ if stats is None:
+ return None
+ else:
+ return DmlStats.from_api_repr(stats)
+
  def _blocking_poll(self, timeout=None, **kwargs):
  self._done_timeout = timeout
  self._transport_timeout = timeout

tests/system/test_client.py

@@ -1521,6 +1521,62 @@ def test_query_statistics(self):
  self.assertGreater(stages_with_inputs, 0)
  self.assertGreater(len(plan), stages_with_inputs)
 
+ def test_dml_statistics(self):
+ table_schema = (
+ bigquery.SchemaField("foo", "STRING"),
+ bigquery.SchemaField("bar", "INTEGER"),
+ )
+
+ dataset_id = _make_dataset_id("bq_system_test")
+ self.temp_dataset(dataset_id)
+ table_id = "{}.{}.test_dml_statistics".format(Config.CLIENT.project, dataset_id)
+
+ # Create the table before loading so that the column order is deterministic.
+ table = helpers.retry_403(Config.CLIENT.create_table)(
+ Table(table_id, schema=table_schema)
+ )
+ self.to_delete.insert(0, table)
+
+ # Insert a few rows and check the stats.
+ sql = f"""
+ INSERT INTO `{table_id}`
+ VALUES ("one", 1), ("two", 2), ("three", 3), ("four", 4);
+ """
+ query_job = Config.CLIENT.query(sql)
+ query_job.result()
+
+ assert query_job.dml_stats is not None
+ assert query_job.dml_stats.inserted_row_count == 4
+ assert query_job.dml_stats.updated_row_count == 0
+ assert query_job.dml_stats.deleted_row_count == 0
+
+ # Update some of the rows.
+ sql = f"""
+ UPDATE `{table_id}`
+ SET bar = bar + 1
+ WHERE bar > 2;
+ """
+ query_job = Config.CLIENT.query(sql)
+ query_job.result()
+
+ assert query_job.dml_stats is not None
+ assert query_job.dml_stats.inserted_row_count == 0
+ assert query_job.dml_stats.updated_row_count == 2
+ assert query_job.dml_stats.deleted_row_count == 0
+
+ # Now delete a few rows and check the stats.
+ sql = f"""
+ DELETE FROM `{table_id}`
+ WHERE foo != "two";
+ """
+ query_job = Config.CLIENT.query(sql)
+ query_job.result()
+
+ assert query_job.dml_stats is not None
+ assert query_job.dml_stats.inserted_row_count == 0
+ assert query_job.dml_stats.updated_row_count == 0
+ assert query_job.dml_stats.deleted_row_count == 3
+
  def test_dbapi_w_standard_sql_types(self):
  for sql, expected in helpers.STANDARD_SQL_EXAMPLES:
  Config.CURSOR.execute(sql)

tests/unit/job/test_query.py

@@ -110,6 +110,24 @@ def _verify_table_definitions(self, job, config):
  self.assertIsNotNone(expected_ec)
  self.assertEqual(found_ec.to_api_repr(), expected_ec)
 
+ def _verify_dml_stats_resource_properties(self, job, resource):
+ query_stats = resource.get("statistics", {}).get("query", {})
+
+ if "dmlStats" in query_stats:
+ resource_dml_stats = query_stats["dmlStats"]
+ job_dml_stats = job.dml_stats
+ assert str(job_dml_stats.inserted_row_count) == resource_dml_stats.get(
+ "insertedRowCount", "0"
+ )
+ assert str(job_dml_stats.updated_row_count) == resource_dml_stats.get(
+ "updatedRowCount", "0"
+ )
+ assert str(job_dml_stats.deleted_row_count) == resource_dml_stats.get(
+ "deletedRowCount", "0"
+ )
+ else:
+ assert job.dml_stats is None
+
  def _verify_configuration_properties(self, job, configuration):
  if "dryRun" in configuration:
  self.assertEqual(job.dry_run, configuration["dryRun"])
@@ -118,6 +136,7 @@ def _verify_configuration_properties(self, job, configuration):
 
  def _verifyResourceProperties(self, job, resource):
  self._verifyReadonlyResourceProperties(job, resource)
+ self._verify_dml_stats_resource_properties(job, resource)
 
  configuration = resource.get("configuration", {})
  self._verify_configuration_properties(job, configuration)
@@ -130,16 +149,19 @@ def _verifyResourceProperties(self, job, resource):
  self._verify_table_definitions(job, query_config)
 
  self.assertEqual(job.query, query_config["query"])
+
  if "createDisposition" in query_config:
  self.assertEqual(job.create_disposition, query_config["createDisposition"])
  else:
  self.assertIsNone(job.create_disposition)
+
  if "defaultDataset" in query_config:
  ds_ref = job.default_dataset
  ds_ref = {"projectId": ds_ref.project, "datasetId": ds_ref.dataset_id}
  self.assertEqual(ds_ref, query_config["defaultDataset"])
  else:
  self.assertIsNone(job.default_dataset)
+
  if "destinationTable" in query_config:
  table = job.destination
  tb_ref = {
@@ -150,14 +172,17 @@ def _verifyResourceProperties(self, job, resource):
  self.assertEqual(tb_ref, query_config["destinationTable"])
  else:
  self.assertIsNone(job.destination)
+
  if "priority" in query_config:
  self.assertEqual(job.priority, query_config["priority"])
  else:
  self.assertIsNone(job.priority)
+
  if "writeDisposition" in query_config:
  self.assertEqual(job.write_disposition, query_config["writeDisposition"])
  else:
  self.assertIsNone(job.write_disposition)
+
  if "destinationEncryptionConfiguration" in query_config:
  self.assertIsNotNone(job.destination_encryption_configuration)
  self.assertEqual(
@@ -166,6 +191,7 @@ def _verifyResourceProperties(self, job, resource):
  )
  else:
  self.assertIsNone(job.destination_encryption_configuration)
+
  if "schemaUpdateOptions" in query_config:
  self.assertEqual(
  job.schema_update_options, query_config["schemaUpdateOptions"]
@@ -190,6 +216,7 @@ def test_ctor_defaults(self):
  self.assertIsNone(job.create_disposition)
  self.assertIsNone(job.default_dataset)
  self.assertIsNone(job.destination)
+ self.assertIsNone(job.dml_stats)
  self.assertIsNone(job.flatten_results)
  self.assertIsNone(job.priority)
  self.assertIsNone(job.use_query_cache)
@@ -278,6 +305,26 @@ def test_from_api_repr_with_encryption(self):
  self.assertIs(job._client, client)
  self._verifyResourceProperties(job, RESOURCE)
 
+ def test_from_api_repr_with_dml_stats(self):
+ self._setUpConstants()
+ client = _make_client(project=self.PROJECT)
+ RESOURCE = {
+ "id": self.JOB_ID,
+ "jobReference": {"projectId": self.PROJECT, "jobId": self.JOB_ID},
+ "configuration": {"query": {"query": self.QUERY}},
+ "statistics": {
+ "query": {
+ "dmlStats": {"insertedRowCount": "15", "updatedRowCount": "2"},
+ },
+ },
+ }
+ klass = self._get_target_class()
+
+ job = klass.from_api_repr(RESOURCE, client=client)
+
+ self.assertIs(job._client, client)
+ self._verifyResourceProperties(job, RESOURCE)
+
  def test_from_api_repr_w_properties(self):
  from google.cloud.bigquery.job import CreateDisposition
  from google.cloud.bigquery.job import SchemaUpdateOption
@@ -815,6 +862,23 @@ def test_estimated_bytes_processed(self):
  query_stats["estimatedBytesProcessed"] = str(est_bytes)
  self.assertEqual(job.estimated_bytes_processed, est_bytes)
 
+ def test_dml_stats(self):
+ from google.cloud.bigquery.job.query import DmlStats
+
+ client = _make_client(project=self.PROJECT)
+ job = self._make_one(self.JOB_ID, self.QUERY, client)
+ assert job.dml_stats is None
+
+ statistics = job._properties["statistics"] = {}
+ assert job.dml_stats is None
+
+ query_stats = statistics["query"] = {}
+ assert job.dml_stats is None
+
+ query_stats["dmlStats"] = {"insertedRowCount": "35"}
+ assert isinstance(job.dml_stats, DmlStats)
+ assert job.dml_stats.inserted_row_count == 35
+
  def test_result(self):
  from google.cloud.bigquery.table import RowIterator
 

tests/unit/job/test_query_stats.py

@@ -15,6 +15,43 @@
 from .helpers import _Base
 
 
+class TestDmlStats:
+ @staticmethod
+ def _get_target_class():
+ from google.cloud.bigquery.job import DmlStats
+
+ return DmlStats
+
+ def _make_one(self, *args, **kw):
+ return self._get_target_class()(*args, **kw)
+
+ def test_ctor_defaults(self):
+ dml_stats = self._make_one()
+ assert dml_stats.inserted_row_count == 0
+ assert dml_stats.deleted_row_count == 0
+ assert dml_stats.updated_row_count == 0
+
+ def test_from_api_repr_partial_stats(self):
+ klass = self._get_target_class()
+ result = klass.from_api_repr({"deletedRowCount": "12"})
+
+ assert isinstance(result, klass)
+ assert result.inserted_row_count == 0
+ assert result.deleted_row_count == 12
+ assert result.updated_row_count == 0
+
+ def test_from_api_repr_full_stats(self):
+ klass = self._get_target_class()
+ result = klass.from_api_repr(
+ {"updatedRowCount": "4", "insertedRowCount": "7", "deletedRowCount": "25"}
+ )
+
+ assert isinstance(result, klass)
+ assert result.inserted_row_count == 7
+ assert result.deleted_row_count == 25
+ assert result.updated_row_count == 4
+
+
 class TestQueryPlanEntryStep(_Base):
  KIND = "KIND"
  SUBSTEPS = ("SUB1", "SUB2")