feat: add additional query stats (#2270)

shollyman · web-flow · commit 7b1b718123af · 2025-09-13T08:02:19.000-07:00
* feat: add additional query stats

This PR adds support for incremental query stats.
diff --git a/google/cloud/bigquery/job/__init__.py b/google/cloud/bigquery/job/__init__.py
@@ -39,6 +39,7 @@
 from google.cloud.bigquery.job.query import QueryPlanEntryStep
 from google.cloud.bigquery.job.query import ScriptOptions
 from google.cloud.bigquery.job.query import TimelineEntry
+from google.cloud.bigquery.job.query import IncrementalResultStats
 from google.cloud.bigquery.enums import Compression
 from google.cloud.bigquery.enums import CreateDisposition
 from google.cloud.bigquery.enums import DestinationFormat
@@ -84,4 +85,5 @@
     "SourceFormat",
     "TransactionInfo",
     "WriteDisposition",
+    "IncrementalResultStats",
 ]
diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py
@@ -197,6 +197,66 @@ def from_api_repr(cls, stats: Dict[str, str]) -> "DmlStats":
         return cls(*args)
 
 
+class IncrementalResultStats:
+    """IncrementalResultStats provides information about incremental query execution."""
+
+    def __init__(self):
+        self._properties = {}
+
+    @classmethod
+    def from_api_repr(cls, resource) -> "IncrementalResultStats":
+        """Factory: construct instance from the JSON repr.
+
+        Args:
+            resource(Dict[str: object]):
+                IncrementalResultStats representation returned from API.
+
+        Returns:
+            google.cloud.bigquery.job.IncrementalResultStats:
+                stats parsed from ``resource``.
+        """
+        entry = cls()
+        entry._properties = resource
+        return entry
+
+    @property
+    def disabled_reason(self):
+        """Optional[string]: Reason why incremental results were not
+        written by the query.
+        """
+        return _helpers._str_or_none(self._properties.get("disabledReason"))
+
+    @property
+    def result_set_last_replace_time(self):
+        """Optional[datetime]: The time at which the result table's contents
+        were completely replaced.  May be absent if no results have been written
+        or the query has completed."""
+        from google.cloud._helpers import _rfc3339_nanos_to_datetime
+
+        value = self._properties.get("resultSetLastReplaceTime")
+        if value:
+            try:
+                return _rfc3339_nanos_to_datetime(value)
+            except ValueError:
+                pass
+        return None
+
+    @property
+    def result_set_last_modify_time(self):
+        """Optional[datetime]: The time at which the result table's contents
+        were modified. May be absent if no results have been written or the
+        query has completed."""
+        from google.cloud._helpers import _rfc3339_nanos_to_datetime
+
+        value = self._properties.get("resultSetLastModifyTime")
+        if value:
+            try:
+                return _rfc3339_nanos_to_datetime(value)
+            except ValueError:
+                pass
+        return None
+
+
 class IndexUnusedReason(typing.NamedTuple):
     """Reason about why no search index was used in the search query (or sub-query).
 
@@ -1339,6 +1399,13 @@ def bi_engine_stats(self) -> Optional[BiEngineStats]:
         else:
             return BiEngineStats.from_api_repr(stats)
 
+    @property
+    def incremental_result_stats(self) -> Optional[IncrementalResultStats]:
+        stats = self._job_statistics().get("incrementalResultStats")
+        if stats is None:
+            return None
+        return IncrementalResultStats.from_api_repr(stats)
+
     def _blocking_poll(self, timeout=None, **kwargs):
         self._done_timeout = timeout
         self._transport_timeout = timeout
diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py
@@ -838,6 +838,23 @@ def test_search_stats(self):
         assert isinstance(job.search_stats, SearchStats)
         assert job.search_stats.mode == "INDEX_USAGE_MODE_UNSPECIFIED"
 
+    def test_incremental_result_stats(self):
+        from google.cloud.bigquery.job.query import IncrementalResultStats
+
+        client = _make_client(project=self.PROJECT)
+        job = self._make_one(self.JOB_ID, self.QUERY, client)
+        assert job.incremental_result_stats is None
+
+        statistics = job._properties["statistics"] = {}
+        assert job.incremental_result_stats is None
+
+        query_stats = statistics["query"] = {}
+        assert job.incremental_result_stats is None
+
+        query_stats["incrementalResultStats"] = {"disabledReason": "BAZ"}
+        assert isinstance(job.incremental_result_stats, IncrementalResultStats)
+        assert job.incremental_result_stats.disabled_reason == "BAZ"
+
     def test_reload_query_results_uses_transport_timeout(self):
         conn = make_connection({})
         client = _make_client(self.PROJECT, connection=conn)
diff --git a/tests/unit/job/test_query_stats.py b/tests/unit/job/test_query_stats.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from .helpers import _Base
+import datetime
 
 
 class TestBiEngineStats:
@@ -520,3 +521,63 @@ def test_from_api_repr_normal(self):
         self.assertEqual(entry.pending_units, self.PENDING_UNITS)
         self.assertEqual(entry.completed_units, self.COMPLETED_UNITS)
         self.assertEqual(entry.slot_millis, self.SLOT_MILLIS)
+
+
+class TestIncrementalResultStats:
+    @staticmethod
+    def _get_target_class():
+        from google.cloud.bigquery.job import IncrementalResultStats
+
+        return IncrementalResultStats
+
+    def _make_one(self, *args, **kw):
+        return self._get_target_class()(*args, **kw)
+
+    def test_ctor_defaults(self):
+        stats = self._make_one()
+        assert stats.disabled_reason is None
+        assert stats.result_set_last_replace_time is None
+        assert stats.result_set_last_modify_time is None
+
+    def test_from_api_repr_partial_stats(self):
+        klass = self._get_target_class()
+        stats = klass.from_api_repr({"disabledReason": "FOO"})
+
+        assert isinstance(stats, klass)
+        assert stats.disabled_reason == "FOO"
+        assert stats.result_set_last_replace_time is None
+        assert stats.result_set_last_modify_time is None
+
+    def test_from_api_repr_full_stats(self):
+        klass = self._get_target_class()
+        stats = klass.from_api_repr(
+            {
+                "disabledReason": "BAR",
+                "resultSetLastReplaceTime": "2025-01-02T03:04:05.06Z",
+                "resultSetLastModifyTime": "2025-02-02T02:02:02.02Z",
+            }
+        )
+
+        assert isinstance(stats, klass)
+        assert stats.disabled_reason == "BAR"
+        assert stats.result_set_last_replace_time == datetime.datetime(
+            2025, 1, 2, 3, 4, 5, 60000, tzinfo=datetime.timezone.utc
+        )
+        assert stats.result_set_last_modify_time == datetime.datetime(
+            2025, 2, 2, 2, 2, 2, 20000, tzinfo=datetime.timezone.utc
+        )
+
+    def test_from_api_repr_invalid_stats(self):
+        klass = self._get_target_class()
+        stats = klass.from_api_repr(
+            {
+                "disabledReason": "BAR",
+                "resultSetLastReplaceTime": "xxx",
+                "resultSetLastModifyTime": "yyy",
+            }
+        )
+
+        assert isinstance(stats, klass)
+        assert stats.disabled_reason == "BAR"
+        assert stats.result_set_last_replace_time is None
+        assert stats.result_set_last_modify_time is None
