feat: add total slot ms to RowIterator

google/cloud/bigquery/_job_helpers.py

@@ -560,6 +560,7 @@ def do_query():
  num_dml_affected_rows=query_results.num_dml_affected_rows,
  query=query,
  total_bytes_processed=query_results.total_bytes_processed,
+ slot_millis=query_results.slot_millis,
  )
 
  if job_retry is not None:

google/cloud/bigquery/client.py

@@ -4144,6 +4144,7 @@ def _list_rows_from_query_results(
  num_dml_affected_rows: Optional[int] = None,
  query: Optional[str] = None,
  total_bytes_processed: Optional[int] = None,
+ slot_millis: Optional[int] = None,
  ) -> RowIterator:
  """List the rows of a completed query.
  See
@@ -4195,6 +4196,8 @@ def _list_rows_from_query_results(
  The query text used.
  total_bytes_processed (Optional[int]):
  total bytes processed from job statistics, if present.
+ slot_millis (Optional[int]):
+ Number of slot ms the user is actually billed for.
 
  Returns:
  google.cloud.bigquery.table.RowIterator:
@@ -4234,6 +4237,7 @@ def _list_rows_from_query_results(
  num_dml_affected_rows=num_dml_affected_rows,
  query=query,
  total_bytes_processed=total_bytes_processed,
+ slot_millis=slot_millis,
  )
  return row_iterator
 

google/cloud/bigquery/job/query.py

@@ -1766,6 +1766,7 @@ def is_job_done():
  num_dml_affected_rows=self._query_results.num_dml_affected_rows,
  query=self.query,
  total_bytes_processed=self.total_bytes_processed,
+ slot_millis=self.slot_millis,
  **list_rows_kwargs,
  )
  rows._preserve_order = _contains_order_by(self.query)

google/cloud/bigquery/query.py

@@ -1282,6 +1282,20 @@ def total_bytes_processed(self):
  if total_bytes_processed is not None:
  return int(total_bytes_processed)
 
+ @property
+ def slot_millis(self):
+ """Total number of slot ms the user is actually billed for.
+
+ See:
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#body.QueryResponse.FIELDS.slot_millis
+
+ Returns:
+ Optional[int]: Count generated on the server (None until set by the server).
+ """
+ slot_millis = self._properties.get("totalSlotMs")
+ if slot_millis is not None:
+ return int(slot_millis)
+
  @property
  def num_dml_affected_rows(self):
  """Total number of rows affected by a DML query.

google/cloud/bigquery/table.py

@@ -1812,6 +1812,7 @@ def __init__(
  num_dml_affected_rows: Optional[int] = None,
  query: Optional[str] = None,
  total_bytes_processed: Optional[int] = None,
+ slot_millis: Optional[int] = None,
  ):
  super(RowIterator, self).__init__(
  client,
@@ -1841,6 +1842,7 @@ def __init__(
  self._num_dml_affected_rows = num_dml_affected_rows
  self._query = query
  self._total_bytes_processed = total_bytes_processed
+ self._slot_millis = slot_millis
 
  @property
  def _billing_project(self) -> Optional[str]:
@@ -1898,6 +1900,11 @@ def total_bytes_processed(self) -> Optional[int]:
  """total bytes processed from job statistics, if present."""
  return self._total_bytes_processed
 
+ @property
+ def slot_millis(self) -> Optional[int]:
+ """Number of slot ms the user is actually billed for."""
+ return self._slot_millis
+
  def _is_almost_completely_cached(self):
  """Check if all results are completely cached.
 

tests/unit/job/test_query.py

@@ -888,6 +888,7 @@ def test_result_reloads_job_state_until_done(self):
  job_resource = self._make_resource(started=True, location="EU")
  job_resource_done = self._make_resource(started=True, ended=True, location="EU")
  job_resource_done["statistics"]["query"]["totalBytesProcessed"] = str(1234)
+ job_resource_done["statistics"]["query"]["totalSlotMs"] = str(5678)
  job_resource_done["configuration"]["query"]["destinationTable"] = {
  "projectId": "dest-project",
  "datasetId": "dest_dataset",
@@ -969,6 +970,7 @@ def test_result_reloads_job_state_until_done(self):
  self.assertEqual(result.total_rows, 1)
  self.assertEqual(result.query, job.query)
  self.assertEqual(result.total_bytes_processed, 1234)
+ self.assertEqual(result.slot_millis, 5678)
 
  query_results_path = f"/projects/{self.PROJECT}/queries/{self.JOB_ID}"
  query_results_call = mock.call(

tests/unit/test_client.py

@@ -5718,6 +5718,7 @@ def test_query_and_wait_defaults(self):
  "rows": [{"f": [{"v": "5552452"}]}],
  "queryId": "job_abcDEF_",
  "totalBytesProcessed": 1234,
+ "totalSlotMs": 5678,
  }
  creds = _make_credentials()
  http = object()
@@ -5735,6 +5736,7 @@ def test_query_and_wait_defaults(self):
  self.assertIsNone(rows.location)
  self.assertEqual(rows.query, query)
  self.assertEqual(rows.total_bytes_processed, 1234)
+ self.assertEqual(rows.slot_millis, 5678)
 
  # Verify the request we send is to jobs.query.
  conn.api_request.assert_called_once()

tests/unit/test_query.py

@@ -2000,6 +2000,22 @@ def test_total_bytes_processed_present_string(self):
  query = self._make_one(resource)
  self.assertEqual(query.total_bytes_processed, 123456)
 
+ def test_slot_millis_missing(self):
+ query = self._make_one(self._make_resource())
+ self.assertIsNone(query.slot_millis)
+
+ def test_slot_millis_present_integer(self):
+ resource = self._make_resource()
+ resource["totalSlotMs"] = 123456
+ query = self._make_one(resource)
+ self.assertEqual(query.slot_millis, 123456)
+
+ def test_slot_millis_present_string(self):
+ resource = self._make_resource()
+ resource["totalSlotMs"] = "123456"
+ query = self._make_one(resource)
+ self.assertEqual(query.slot_millis, 123456)
+
  def test_num_dml_affected_rows_missing(self):
  query = self._make_one(self._make_resource())
  self.assertIsNone(query.num_dml_affected_rows)