feat: add total slot ms to RowIterator (#2233)

* feat: add total slot ms to RowIterator

* format fix

---------

Co-authored-by: Chalmer Lowe <[email protected]>

Author: Genesis929

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)