upscope max_slots to jobconfig, update fastpath eval

Author: shollyman

google/cloud/bigquery/_job_helpers.py

@@ -658,6 +658,8 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
         "requestId",
         "createSession",
         "writeIncrementalResults",
+        "jobTimeoutMs",
+        "reservation",
         "maxSlots",
     }
 

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.

google/cloud/bigquery/job/query.py

@@ -689,28 +689,6 @@ def write_incremental_results(self) -> Optional[bool]:
     def write_incremental_results(self, value):
         self._set_sub_prop("writeIncrementalResults", value)
 
-
-    @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._get_sub_prop("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):
-        self._set_sub_prop("maxSlots", value)
-
     @property
     def table_definitions(self):
         """Dict[str, google.cloud.bigquery.external_config.ExternalConfig]:

tests/unit/job/test_query_config.py

@@ -177,7 +177,6 @@ def test_max_slots(self):
         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)

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,9 +1061,19 @@ 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),
-            20,
+            True,
             id="max_slots",
         ),
     ),