Skip to content

Commit d08b253

Browse files
committed
upscope max_slots to jobconfig, update fastpath eval
1 parent f89c551 commit d08b253

File tree

5 files changed

+57
-24
lines changed

5 files changed

+57
-24
lines changed

google/cloud/bigquery/_job_helpers.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,6 +658,8 @@ def _supported_by_jobs_query(request_body: Dict[str, Any]) -> bool:
658658
"requestId",
659659
"createSession",
660660
"writeIncrementalResults",
661+
"jobTimeoutMs",
662+
"reservation",
661663
"maxSlots",
662664
}
663665

google/cloud/bigquery/job/base.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,37 @@ def job_timeout_ms(self, value):
224224
else:
225225
self._properties.pop("jobTimeoutMs", None)
226226

227+
@property
228+
def max_slots(self) -> Optional[int]:
229+
"""The maximum rate of slot consumption to allow for this job.
230+
231+
If set, the number of slots used to execute the job will be throttled
232+
to try and keep its slot consumption below the requested rate.
233+
This feature is not generally available.
234+
"""
235+
236+
max_slots = self._properties.get("maxSlots")
237+
if max_slots is not None:
238+
if isinstance(max_slots, str):
239+
return int(max_slots)
240+
if isinstance(max_slots, int):
241+
return max_slots
242+
return None
243+
244+
@max_slots.setter
245+
def max_slots(self, value):
246+
try:
247+
value = _int_or_none(value)
248+
except ValueError as err:
249+
raise ValueError("Pass an int for max slots, e.g. 100").with_traceback(
250+
err.__traceback__
251+
)
252+
253+
if value is not None:
254+
self._properties["maxSlots"] = str(value)
255+
else:
256+
self._properties.pop("maxSlots", None)
257+
227258
@property
228259
def reservation(self):
229260
"""str: Optional. The reservation that job would use.

google/cloud/bigquery/job/query.py

Lines changed: 0 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -689,28 +689,6 @@ def write_incremental_results(self) -> Optional[bool]:
689689
def write_incremental_results(self, value):
690690
self._set_sub_prop("writeIncrementalResults", value)
691691

692-
693-
@property
694-
def max_slots(self) -> Optional[int]:
695-
"""The maximum rate of slot consumption to allow for this job.
696-
697-
If set, the number of slots used to execute the job will be throttled
698-
to try and keep its slot consumption below the requested rate.
699-
700-
This feature is not generally available.
701-
"""
702-
max_slots = self._get_sub_prop("maxSlots")
703-
if max_slots is not None:
704-
if isinstance(max_slots, str):
705-
return int(max_slots)
706-
if isinstance(max_slots, int):
707-
return max_slots
708-
return None
709-
710-
@max_slots.setter
711-
def max_slots(self, value):
712-
self._set_sub_prop("maxSlots", value)
713-
714692
@property
715693
def table_definitions(self):
716694
"""Dict[str, google.cloud.bigquery.external_config.ExternalConfig]:

tests/unit/job/test_query_config.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,6 @@ def test_max_slots(self):
177177
config.max_slots = 99
178178
self.assertEqual(config.max_slots, 99)
179179

180-
181180
def test_create_session(self):
182181
config = self._get_target_class()()
183182
self.assertIsNone(config.create_session)

tests/unit/test__job_helpers.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,19 @@ def make_query_response(
200200
make_query_request({"writeIncrementalResults": True}),
201201
id="job_config-with-incremental-results",
202202
),
203+
pytest.param(
204+
job_query.QueryJobConfig(
205+
reservation="foo",
206+
max_slots=100,
207+
),
208+
make_query_request(
209+
{
210+
"maxSlots": "100",
211+
"reservation": "foo",
212+
}
213+
),
214+
id="job_config-with-reservation-and-slots",
215+
),
203216
),
204217
)
205218
def test__to_query_request(job_config, expected):
@@ -1048,9 +1061,19 @@ def test_make_job_id_w_job_id_overrides_prefix():
10481061
True,
10491062
id="write_incremental_results",
10501063
),
1064+
pytest.param(
1065+
job_query.QueryJobConfig(job_timeout_ms=1000),
1066+
True,
1067+
id="job_timeout_ms",
1068+
),
1069+
pytest.param(
1070+
job_query.QueryJobConfig(reservation="foo"),
1071+
True,
1072+
id="reservation",
1073+
),
10511074
pytest.param(
10521075
job_query.QueryJobConfig(max_slots=20),
1053-
20,
1076+
True,
10541077
id="max_slots",
10551078
),
10561079
),

0 commit comments

Comments
 (0)