这是indexloc提供的服务,不要输入任何密码
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions google/cloud/bigquery/_job_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
31 changes: 31 additions & 0 deletions google/cloud/bigquery/job/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#QUESTION

What is the impetus for .pop()?

If value is None, delete the key: value pair entirely?
I know we also did it with jobTimeoutMs and I don't get it there either.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

#QUESTION

What is the impetus for .pop()?

If value is None, delete the key: value pair entirely? I know we also did it with jobTimeoutMs and I don't get it there either.

It's the easier one liner. We either need to check for the key existence explicitly, or wrap the delete in a try...except block to handle the cases where you're setting the None value and the key doesn't yet exist in the dict. Happy to swap to one of these if you prefer.


@property
def reservation(self):
"""str: Optional. The reservation that job would use.
Expand Down
41 changes: 41 additions & 0 deletions tests/unit/job/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
5 changes: 5 additions & 0 deletions tests/unit/job/test_query_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
28 changes: 28 additions & 0 deletions tests/unit/test__job_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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(
Expand Down