mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
test(integration): recognize cached statements in database relay barriers (#50651)
Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
eae60c0bc9
commit
d10de278c1
3 changed files with 142 additions and 9 deletions
|
|
@ -91,7 +91,6 @@ class HeldStatementRelay:
|
|||
self._upstream_port: Final = upstream_port
|
||||
self._trigger: Final = trigger
|
||||
self._after: Final = after
|
||||
self._marker_length: Final = max(len(trigger), len(after or b""))
|
||||
self._server: Final = LoopbackServer(self.port, self._serve)
|
||||
self._loop: Final = self._server.loop
|
||||
self._released: Final = asyncio.Event()
|
||||
|
|
@ -115,27 +114,25 @@ class HeldStatementRelay:
|
|||
self.release()
|
||||
self._server.stop()
|
||||
|
||||
def _holds(self, window: bytes) -> bool:
|
||||
def _holds(self, markers: frozenset[bytes]) -> bool:
|
||||
if not self._armed.is_set() or self.held.is_set():
|
||||
return False
|
||||
if self._after is not None and not self._after_seen.is_set():
|
||||
if self._after in window:
|
||||
if self._after in markers:
|
||||
self._after_seen.set()
|
||||
return False
|
||||
return self._trigger in window
|
||||
return self._trigger in markers
|
||||
|
||||
async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None:
|
||||
server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port)
|
||||
|
||||
async def forward(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, inspect: bool) -> None:
|
||||
tail = b"" # rebind-ok: carries the previous read's end so a trigger split across reads still matches
|
||||
scanner: Final = StatementMarkerScanner((self._trigger,) + ((self._after,) if self._after else ()))
|
||||
try:
|
||||
while chunk := await reader.read(65536):
|
||||
window: Final = tail + chunk
|
||||
if inspect and self._holds(window):
|
||||
if inspect and any(self._holds(markers) for markers in scanner.feed(chunk)):
|
||||
self.held.set()
|
||||
await self._released.wait()
|
||||
tail = window[-(self._marker_length - 1) :]
|
||||
writer.write(chunk)
|
||||
await writer.drain()
|
||||
except (ConnectionError, asyncio.IncompleteReadError):
|
||||
|
|
@ -149,6 +146,47 @@ class HeldStatementRelay:
|
|||
)
|
||||
|
||||
|
||||
class StatementMarkerScanner:
|
||||
def __init__(self, markers: tuple[bytes, ...]) -> None:
|
||||
self._markers: Final = markers
|
||||
self._statements: Final[dict[bytes, frozenset[bytes]]] = {}
|
||||
self._buffer = b""
|
||||
self._startup = True
|
||||
|
||||
def feed(self, chunk: bytes) -> Generator[frozenset[bytes]]:
|
||||
self._buffer += chunk
|
||||
while len(self._buffer) >= (4 if self._startup else 5):
|
||||
offset: Final = 0 if self._startup else 1
|
||||
length: Final = int.from_bytes(self._buffer[offset : offset + 4], "big")
|
||||
if length < 4:
|
||||
raise ValueError("Invalid PostgreSQL message length")
|
||||
if len(self._buffer) < length + offset:
|
||||
return
|
||||
message: Final = self._buffer[: length + offset]
|
||||
self._buffer = self._buffer[length + offset :]
|
||||
if self._startup:
|
||||
self._startup = message[4:8] == b"\x04\xd2\x16/"
|
||||
continue
|
||||
payload: Final = message[5:]
|
||||
if message[:1] == b"P":
|
||||
name, _, query = payload.partition(b"\x00")
|
||||
matches: Final = self._matches(query.split(b"\x00", 1)[0])
|
||||
if name:
|
||||
self._statements[name] = matches
|
||||
yield matches
|
||||
elif message[:1] == b"B":
|
||||
fields: Final = payload.split(b"\x00", 2)
|
||||
prepared: Final = self._statements.get(fields[1], frozenset()) if len(fields) > 1 else frozenset()
|
||||
yield prepared | self._matches(payload)
|
||||
elif message[:1] == b"Q":
|
||||
yield self._matches(payload)
|
||||
elif message[:1] == b"C" and payload[:1] == b"S":
|
||||
self._statements.pop(payload[1:].split(b"\x00", 1)[0], None)
|
||||
|
||||
def _matches(self, query: bytes) -> frozenset[bytes]:
|
||||
return frozenset(marker for marker in self._markers if marker in query)
|
||||
|
||||
|
||||
class TriggerScanner:
|
||||
def __init__(self, trigger: bytes) -> None:
|
||||
self._trigger: Final = trigger
|
||||
|
|
|
|||
|
|
@ -219,6 +219,49 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite
|
|||
await reader.read()
|
||||
|
||||
|
||||
async def _echo_protocol(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
||||
while chunk := await reader.read(65536):
|
||||
writer.write(chunk)
|
||||
await writer.drain()
|
||||
|
||||
|
||||
def _protocol_message(kind: bytes, payload: bytes) -> bytes:
|
||||
return kind + (len(payload) + 4).to_bytes(4, "big") + payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prepared_trigger", (False, True))
|
||||
def test_relay_holds_cached_statements_parsed_before_arming(prepared_trigger: bool) -> None:
|
||||
upstream_port: Final = allocate_port()
|
||||
upstream: Final = LoopbackServer(upstream_port, _echo_protocol)
|
||||
relay: Final = HeldStatementRelay("127.0.0.1", upstream_port, b"trigger", after=b"snapshot", armed=False)
|
||||
startup: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00"
|
||||
snapshot: Final = _protocol_message(b"P", b"s1\x00SELECT snapshot\x00\x00\x00")
|
||||
trigger: Final = _protocol_message(b"P", b"s2\x00SELECT trigger\x00\x00\x00")
|
||||
bind_snapshot: Final = _protocol_message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00")
|
||||
bind_trigger: Final = _protocol_message(b"B", b"\x00s2\x00\x00\x00\x00\x00\x00\x00")
|
||||
upstream.start()
|
||||
try:
|
||||
relay.start()
|
||||
with (
|
||||
socket.create_connection(("127.0.0.1", relay.port), timeout=5) as client,
|
||||
client.makefile("rb") as replies,
|
||||
):
|
||||
before: Final = startup + snapshot + bind_snapshot + (trigger if prepared_trigger else b"")
|
||||
client.sendall(before)
|
||||
assert replies.read(len(before)) == before
|
||||
relay.arm()
|
||||
client.sendall(bind_snapshot)
|
||||
assert replies.read(len(bind_snapshot)) == bind_snapshot
|
||||
held: Final = bind_trigger if prepared_trigger else _protocol_message(b"Q", b"SELECT trigger\x00")
|
||||
client.sendall(held)
|
||||
assert relay.held.wait(timeout=5), "cached snapshot did not open the relay barrier"
|
||||
relay.release()
|
||||
assert replies.read(len(held)) == held
|
||||
finally:
|
||||
relay.stop()
|
||||
upstream.stop()
|
||||
|
||||
|
||||
def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None:
|
||||
port: Final = allocate_port()
|
||||
with socket.socket() as occupied:
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from typing import Final
|
|||
|
||||
import pytest
|
||||
|
||||
from tests.integration._support.database_relay import TriggerScanner
|
||||
from tests.integration._support.database_relay import StatementMarkerScanner, TriggerScanner
|
||||
|
||||
TRIGGER: Final = b'SELECT "startTime" FROM "LiteLLM_SpendLogs"'
|
||||
|
||||
|
|
@ -26,3 +26,55 @@ def test_trigger_scanner_reports_a_match_once() -> None:
|
|||
scanner: Final = TriggerScanner(TRIGGER)
|
||||
assert scanner.feed(b"x" + TRIGGER + b"y")
|
||||
assert not scanner.feed(b"z")
|
||||
|
||||
|
||||
STARTUP: Final = b"\x00\x00\x00\x08\x00\x03\x00\x00"
|
||||
|
||||
|
||||
def _message(kind: bytes, payload: bytes) -> bytes:
|
||||
return kind + (len(payload) + 4).to_bytes(4, "big") + payload
|
||||
|
||||
|
||||
PARSE: Final = _message(b"P", b"s1\x00" + TRIGGER + b"\x00\x00\x00")
|
||||
BIND: Final = _message(b"B", b"\x00s1\x00\x00\x00\x00\x00\x00\x00")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("split_at", range(1, len(STARTUP + PARSE + BIND)))
|
||||
def test_statement_markers_survive_every_protocol_frame_split(split_at: int) -> None:
|
||||
scanner: Final = StatementMarkerScanner((TRIGGER,))
|
||||
stream: Final = STARTUP + PARSE + BIND
|
||||
assert (*scanner.feed(stream[:split_at]), *scanner.feed(stream[split_at:])) == (
|
||||
frozenset({TRIGGER}),
|
||||
frozenset({TRIGGER}),
|
||||
)
|
||||
|
||||
|
||||
def test_statement_markers_are_local_to_a_connection_and_forget_closed_statements() -> None:
|
||||
scanner: Final = StatementMarkerScanner((TRIGGER,))
|
||||
other: Final = StatementMarkerScanner((TRIGGER,))
|
||||
assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),)
|
||||
assert tuple(other.feed(STARTUP + BIND)) == (frozenset(),)
|
||||
assert tuple(scanner.feed(_message(b"C", b"Ss1\x00") + BIND)) == (frozenset(),)
|
||||
|
||||
|
||||
def test_replaced_and_unnamed_statements_do_not_reuse_old_markers() -> None:
|
||||
scanner: Final = StatementMarkerScanner((TRIGGER,))
|
||||
assert tuple(scanner.feed(STARTUP + PARSE)) == (frozenset({TRIGGER}),)
|
||||
replacement: Final = _message(b"P", b"s1\x00SELECT 1\x00\x00\x00")
|
||||
unnamed: Final = _message(b"P", b"\x00" + TRIGGER + b"\x00\x00\x00")
|
||||
assert tuple(scanner.feed(replacement + BIND + unnamed)) == (frozenset(), frozenset(), frozenset({TRIGGER}))
|
||||
|
||||
|
||||
def test_statement_scanner_rejects_invalid_frame_lengths() -> None:
|
||||
scanner: Final = StatementMarkerScanner((TRIGGER,))
|
||||
assert tuple(scanner.feed(STARTUP)) == ()
|
||||
with pytest.raises(ValueError, match="Invalid PostgreSQL message length"):
|
||||
tuple(scanner.feed(b"P\x00\x00\x00\x03"))
|
||||
|
||||
|
||||
def test_statement_scanner_matches_bound_parameter_markers() -> None:
|
||||
scanner: Final = StatementMarkerScanner((TRIGGER,))
|
||||
bind: Final = _message(
|
||||
b"B", b"\x00s2\x00\x00\x00\x00\x01" + len(TRIGGER).to_bytes(4, "big") + TRIGGER + b"\x00\x00"
|
||||
)
|
||||
assert tuple(scanner.feed(STARTUP + bind)) == (frozenset({TRIGGER}),)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue