From d10de278c1bb65a271c032e4e17359bb3f0c64f9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 9 Nov 2026 06:58:46 -0800 Subject: [PATCH] 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> --- tests/integration/_support/database_relay.py | 54 ++++++++++++++++--- .../database/test_harness_lifecycle.py | 43 +++++++++++++++ .../test_database_relay.py | 54 ++++++++++++++++++- 3 files changed, 142 insertions(+), 9 deletions(-) diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 0ac1e1d2521..0ac060dc867 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -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 diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index d600a92c5a2..4cb01dfeef6 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -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: diff --git a/tests/unit/integration_support/test_database_relay.py b/tests/unit/integration_support/test_database_relay.py index 98501f84058..de2bf1f5d2b 100644 --- a/tests/unit/integration_support/test_database_relay.py +++ b/tests/unit/integration_support/test_database_relay.py @@ -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}),)