Repository: googleapis/python-bigquery (python)

Issue:
totalBytesProcessed is not exposed from RowIterator
RowIterator resulting from `query_and_wait()` for the new short query mode exposes most of its internal state via properties (cf. e.g. `num_dml_affected_rows`) but does not provide a `total_bytes_processed` property to access `self.totalBytesProcessed`. It would be nice if as many as possible of the same attributes were present to make transitioning to `query_and_wait()` from `query()` as smooth as possible.

I'm sure there are other attributes too but this is the one that I ran into first.

Pull request:
feat: add query text and total bytes processed to RowIterator
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)

Fixes b/400961399
Fixes #1996
🦕


Patch:
diff --git a/google/cloud/bigquery/_job_helpers.py b/google/cloud/bigquery/_job_helpers.py
index b028cd35..a8373c35 100644
--- a/google/cloud/bigquery/_job_helpers.py
+++ b/google/cloud/bigquery/_job_helpers.py
@@ -526,6 +526,8 @@ def query_and_wait(
             query_id=query_results.query_id,
             project=query_results.project,
             num_dml_affected_rows=query_results.num_dml_affected_rows,
+            query=query,
+            total_bytes_processed=query_results.total_bytes_processed,
         )
 
     if job_retry is not None:
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
index 03ded93b..8bbdd6c3 100644
--- a/google/cloud/bigquery/client.py
+++ b/google/cloud/bigquery/client.py
@@ -4081,6 +4081,8 @@ class Client(ClientWithProject):
         query_id: Optional[str] = None,
         first_page_response: Optional[Dict[str, Any]] = None,
         num_dml_affected_rows: Optional[int] = None,
+        query: Optional[str] = None,
+        total_bytes_processed: Optional[int] = None,
     ) -> RowIterator:
         """List the rows of a completed query.
         See
@@ -4128,6 +4130,10 @@ class Client(ClientWithProject):
             num_dml_affected_rows (Optional[int]):
                 If this RowIterator is the result of a DML query, the number of
                 rows that were affected.
+            query (Optional[str]):
+                The query text used.
+            total_bytes_processed (Optinal[int]):
+                total bytes processed from job statistics, if present.
 
         Returns:
             google.cloud.bigquery.table.RowIterator:
@@ -4165,6 +4171,8 @@ class Client(ClientWithProject):
             query_id=query_id,
             first_page_response=first_page_response,
             num_dml_affected_rows=num_dml_affected_rows,
+            query=query,
+            total_bytes_processed=total_bytes_processed,
         )
         return row_iterator
 
diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
index ca2448ea..a27c1053 100644
--- a/google/cloud/bigquery/job/query.py
+++ b/google/cloud/bigquery/job/query.py
@@ -1741,6 +1741,8 @@ class QueryJob(_AsyncJob):
             query_id=self.query_id,
             first_page_response=first_page_response,
             num_dml_affected_rows=self._query_results.num_dml_affected_rows,
+            query=self.query,
+            total_bytes_processed=self.total_bytes_processed,
             **list_rows_kwargs,
         )
         rows._preserve_order = _contains_order_by(self.query)
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index 9950b1a5..4d79d60d 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -1760,6 +1760,10 @@ class RowIterator(HTTPIterator):
         first_page_response (Optional[dict]):
             API response for the first page of results. These are returned when
             the first page is requested.
+        query (Optional[str]):
+            The query text used.
+        total_bytes_processed (Optinal[int]):
+            total bytes processed from job statistics, if present.
     """
 
     def __init__(
@@ -1781,6 +1785,8 @@ class RowIterator(HTTPIterator):
         query_id: Optional[str] = None,
         project: Optional[str] = None,
         num_dml_affected_rows: Optional[int] = None,
+        query: Optional[str] = None,
+        total_bytes_processed: Optional[int] = None,
     ):
         super(RowIterator, self).__init__(
             client,
@@ -1808,6 +1814,8 @@ class RowIterator(HTTPIterator):
         self._query_id = query_id
         self._project = project
         self._num_dml_affected_rows = num_dml_affected_rows
+        self._query = query
+        self._total_bytes_processed = total_bytes_processed
 
     @property
     def _billing_project(self) -> Optional[str]:
@@ -1855,6 +1863,16 @@ class RowIterator(HTTPIterator):
         """
         return self._query_id
 
+    @property
+    def query(self) -> Optional[str]:
+        """The query text used."""
+        return self._query
+
+    @property
+    def total_bytes_processed(self) -> Optional[int]:
+        """total bytes processed from job statistics, if present."""
+        return self._total_bytes_processed
+
     def _is_almost_completely_cached(self):
         """Check if all results are completely cached.
 


Test patch:
diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py
index 4bbd31c7..1df65279 100644
--- a/tests/unit/job/test_query.py
+++ b/tests/unit/job/test_query.py
@@ -887,6 +887,7 @@ class TestQueryJob(_Base):
         }
         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["configuration"]["query"]["destinationTable"] = {
             "projectId": "dest-project",
             "datasetId": "dest_dataset",
@@ -966,6 +967,8 @@ class TestQueryJob(_Base):
         # Test that the total_rows property has changed during iteration, based
         # on the response from tabledata.list.
         self.assertEqual(result.total_rows, 1)
+        self.assertEqual(result.query, job.query)
+        self.assertEqual(result.total_bytes_processed, 1234)
 
         query_results_path = f"/projects/{self.PROJECT}/queries/{self.JOB_ID}"
         query_results_call = mock.call(
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 4f13d6ec..34ef680d 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -5517,6 +5517,7 @@ class TestClient(unittest.TestCase):
             "totalRows": "1",
             "rows": [{"f": [{"v": "5552452"}]}],
             "queryId": "job_abcDEF_",
+            "totalBytesProcessed": 1234,
         }
         creds = _make_credentials()
         http = object()
@@ -5532,6 +5533,8 @@ class TestClient(unittest.TestCase):
         self.assertIsNone(rows.job_id)
         self.assertIsNone(rows.project)
         self.assertIsNone(rows.location)
+        self.assertEqual(rows.query, query)
+        self.assertEqual(rows.total_bytes_processed, 1234)
 
         # Verify the request we send is to jobs.query.
         conn.api_request.assert_called_once()