feat: updates to fastpath query execution (#2268)

shollyman · web-flow · commit ef2740a15819 · 2025-08-26T15:52:30.000Z
This PR updates query handling to allow base config properties like job timeout, reservation, and a preview max slots field to leverage the faster path (e.g. using jobs.query rather than jobs.insert).
diff --git a/google/cloud/bigquery/_job_helpers.py b/google/cloud/bigquery/_job_helpers.py
@@ -658,6 +658,9 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
         "requestId",
         "createSession",
         "writeIncrementalResults",
+        "jobTimeoutMs",
+        "reservation",
+        "maxSlots",
     }
 
     unsupported_keys = request_keys - keys_allowlist
diff --git a/google/cloud/bigquery/job/base.py b/google/cloud/bigquery/job/base.py
@@ -224,6 +224,37 @@ def job_timeout_ms(self, value):
         else:
             self._properties.pop("jobTimeoutMs", None)
 
+    @property
+    def max_slots(self) -> Optional[int]:
+        """The maximum rate of slot consumption to allow for this job.
+
+        If set, the number of slots used to execute the job will be throttled
+        to try and keep its slot consumption below the requested rate.
+        This feature is not generally available.
+        """
+
+        max_slots = self._properties.get("maxSlots")
+        if max_slots is not None:
+            if isinstance(max_slots, str):
+                return int(max_slots)
+            if isinstance(max_slots, int):
+                return max_slots
+        return None
+
+    @max_slots.setter
+    def max_slots(self, value):
+        try:
+            value = _int_or_none(value)
+        except ValueError as err:
+            raise ValueError("Pass an int for max slots, e.g. 100").with_traceback(
+                err.__traceback__
+            )
+
+        if value is not None:
+            self._properties["maxSlots"] = str(value)
+        else:
+            self._properties.pop("maxSlots", None)
+
     @property
     def reservation(self):
         """str: Optional. The reservation that job would use.
diff --git a/tests/unit/job/test_base.py b/tests/unit/job/test_base.py
@@ -1276,3 +1276,44 @@ def test_reservation_setter(self):
         job_config = self._make_one()
         job_config.reservation = "foo"
         self.assertEqual(job_config._properties["reservation"], "foo")
+
+    def test_max_slots_miss(self):
+        job_config = self._make_one()
+        self.assertEqual(job_config.max_slots, None)
+
+    def test_max_slots_set_and_clear(self):
+        job_config = self._make_one()
+        job_config.max_slots = 14
+        self.assertEqual(job_config.max_slots, 14)
+        job_config.max_slots = None
+        self.assertEqual(job_config.max_slots, None)
+
+    def test_max_slots_hit_str(self):
+        job_config = self._make_one()
+        job_config._properties["maxSlots"] = "4"
+        self.assertEqual(job_config.max_slots, 4)
+
+    def test_max_slots_hit_int(self):
+        job_config = self._make_one()
+        job_config._properties["maxSlots"] = int(3)
+        self.assertEqual(job_config.max_slots, 3)
+
+    def test_max_slots_hit_invalid(self):
+        job_config = self._make_one()
+        job_config._properties["maxSlots"] = object()
+        self.assertEqual(job_config.max_slots, None)
+
+    def test_max_slots_update_in_place(self):
+        job_config = self._make_one()
+        job_config.max_slots = 45  # update in place
+        self.assertEqual(job_config.max_slots, 45)
+
+    def test_max_slots_setter_invalid(self):
+        job_config = self._make_one()
+        with self.assertRaises(ValueError):
+            job_config.max_slots = "foo"
+
+    def test_max_slots_setter(self):
+        job_config = self._make_one()
+        job_config.max_slots = 123
+        self.assertEqual(job_config._properties["maxSlots"], "123")
diff --git a/tests/unit/job/test_query_config.py b/tests/unit/job/test_query_config.py
@@ -172,6 +172,11 @@ def test_incremental_results(self):
         config.write_incremental_results = True
         self.assertEqual(config.write_incremental_results, True)
 
+    def test_max_slots(self):
+        config = self._get_target_class()()
+        config.max_slots = 99
+        self.assertEqual(config.max_slots, 99)
+
     def test_create_session(self):
         config = self._get_target_class()()
         self.assertIsNone(config.create_session)
diff --git a/tests/unit/test__job_helpers.py b/tests/unit/test__job_helpers.py
@@ -200,6 +200,19 @@ def make_query_response(
             make_query_request({"writeIncrementalResults": True}),
             id="job_config-with-incremental-results",
         ),
+        pytest.param(
+            job_query.QueryJobConfig(
+                reservation="foo",
+                max_slots=100,
+            ),
+            make_query_request(
+                {
+                    "maxSlots": "100",
+                    "reservation": "foo",
+                }
+            ),
+            id="job_config-with-reservation-and-slots",
+        ),
     ),
 )
 def test__to_query_request(job_config, expected):
@@ -1048,6 +1061,21 @@ def test_make_job_id_w_job_id_overrides_prefix():
             True,
             id="write_incremental_results",
         ),
+        pytest.param(
+            job_query.QueryJobConfig(job_timeout_ms=1000),
+            True,
+            id="job_timeout_ms",
+        ),
+        pytest.param(
+            job_query.QueryJobConfig(reservation="foo"),
+            True,
+            id="reservation",
+        ),
+        pytest.param(
+            job_query.QueryJobConfig(max_slots=20),
+            True,
+            id="max_slots",
+        ),
     ),
 )
 def test_supported_by_jobs_query_from_queryjobconfig(
