Skip to content

Commit b9a760a

Browse files
authored
Check that advisory lock prefixes fit in 32 bits (#16)
We only reserve 32 bits for an advisory lock prefix, and since Python doesn't have a specific `int32` type, do a preflight check on the client to make sure that a given prefix is valid and will fit.
1 parent 9f2c368 commit b9a760a

File tree

2 files changed

+33
-2
lines changed

2 files changed

+33
-2
lines changed

src/riverqueue/client.py

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,9 @@ def __init__(
4949
self, driver: AsyncDriverProtocol, advisory_lock_prefix: Optional[int] = None
5050
):
5151
self.driver = driver
52-
self.advisory_lock_prefix = advisory_lock_prefix
52+
self.advisory_lock_prefix = _check_advisory_lock_prefix_bounds(
53+
advisory_lock_prefix
54+
)
5355

5456
async def insert(
5557
self, args: Args, insert_opts: Optional[InsertOpts] = None
@@ -122,7 +124,9 @@ def __init__(
122124
self, driver: DriverProtocol, advisory_lock_prefix: Optional[int] = None
123125
):
124126
self.driver = driver
125-
self.advisory_lock_prefix = advisory_lock_prefix
127+
self.advisory_lock_prefix = _check_advisory_lock_prefix_bounds(
128+
advisory_lock_prefix
129+
)
126130

127131
def insert(
128132
self, args: Args, insert_opts: Optional[InsertOpts] = None
@@ -247,6 +251,17 @@ def _build_unique_get_params_and_lock_key(
247251
return (get_params, _uint64_to_int64(lock_key))
248252

249253

254+
def _check_advisory_lock_prefix_bounds(
255+
advisory_lock_prefix: Optional[int],
256+
) -> Optional[int]:
257+
if advisory_lock_prefix:
258+
print("in_bytes", advisory_lock_prefix.to_bytes(4))
259+
# We only reserve 4 bytes for the prefix, so make sure the given one
260+
# properly fits. This will error in case that's not the case.
261+
advisory_lock_prefix.to_bytes(4)
262+
return advisory_lock_prefix
263+
264+
250265
def _make_insert_params(
251266
args: Args,
252267
insert_opts: InsertOpts,

tests/client_test.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,3 +147,19 @@ def test_insert_with_unique_opts_by_state(client, mock_exec):
147147
# Check that the UniqueOpts were correctly processed
148148
call_args = mock_exec.job_insert.call_args[0][0]
149149
assert call_args.kind == "simple"
150+
151+
152+
def test_check_advisory_lock_prefix_bounds():
153+
Client(mock_driver, advisory_lock_prefix=123)
154+
155+
with pytest.raises(OverflowError) as ex:
156+
Client(mock_driver, advisory_lock_prefix=-1)
157+
assert "can't convert negative int to unsigned" == str(ex.value)
158+
159+
# 2^32-1 is 0xffffffff (1s for 32 bits) which fits
160+
Client(mock_driver, advisory_lock_prefix=2**32 - 1)
161+
162+
# 2^32 is 0x100000000, which does not
163+
with pytest.raises(OverflowError) as ex:
164+
Client(mock_driver, advisory_lock_prefix=2**32)
165+
assert "int too big to convert" == str(ex.value)

0 commit comments

Comments
 (0)