From 879c74f34dbebf41c2fcbf8c6b7f99806f48f542 Mon Sep 17 00:00:00 2001 From: Dainius Jocas Date: Sun, 20 Sep 2026 15:34:42 +0200 Subject: [PATCH 1/2] flag to store a deep copy of the request body in the response --- tests/unit/test_application.py | 128 +++++++++++++++++++++++++++++++++ vespa/application.py | 13 ++++ vespa/io.py | 2 + 3 files changed, 143 insertions(+) diff --git a/tests/unit/test_application.py b/tests/unit/test_application.py index 7d0a3ff4a..295322b8b 100644 --- a/tests/unit/test_application.py +++ b/tests/unit/test_application.py @@ -1694,3 +1694,131 @@ def stream(*args, **kwargs): assert list(lines) == ["event: token", 'data: {"token":"V"}'] assert client.closed client.close.assert_called_once() + + +class TestQueryKeepRequestBody(unittest.TestCase): + """Tests for the `keep_request_body` flag on sync query methods.""" + + def _mock_client(self, MockClient): + mock_client_instance = Mock() + MockClient.return_value = mock_client_instance + mock_client_instance.close = Mock() + + status_response = create_mock_httpr_response(status_code=200) + search_response = create_mock_httpr_response( + status_code=200, text="{}", url="http://localhost:8080/search/" + ) + mock_client_instance.get.return_value = status_response + mock_client_instance.post.return_value = search_response + return mock_client_instance + + @patch("vespa.application.httpr.Client") + def test_request_body_none_by_default(self, MockClient): + self._mock_client(MockClient) + + app = Vespa(url="http://localhost", port=8080) + r = app.query(body={"yql": "select * from sources * where true"}) + + self.assertIsNone(r.request_body) + + @patch("vespa.application.httpr.Client") + def test_request_body_stored_when_flag_true(self, MockClient): + self._mock_client(MockClient) + body = {"yql": "select * from sources * where true", "nested": {"a": 1}} + + app = Vespa(url="http://localhost", port=8080) + r = app.query(body=body, keep_request_body=True) + + self.assertEqual(r.request_body, body) + + @patch("vespa.application.httpr.Client") + def test_request_body_is_deep_copy(self, MockClient): + self._mock_client(MockClient) + body = {"yql": "select * from sources * where true", "nested": {"a": 1}} + + app = Vespa(url="http://localhost", port=8080) + r = app.query(body=body, keep_request_body=True) + + self.assertIsNot(r.request_body, body) + self.assertIsNot(r.request_body["nested"], body["nested"]) + + # Mutating the original body after the call must not affect the stored copy. + body["nested"]["a"] = 999 + self.assertEqual(r.request_body["nested"]["a"], 1) + + @patch("vespa.application.httpr.Client") + def test_keep_request_body_not_sent_as_query_param(self, MockClient): + mock_client_instance = self._mock_client(MockClient) + + app = Vespa(url="http://localhost", port=8080) + app.query( + body={"yql": "select * from sources * where true"}, + keep_request_body=True, + ) + + _, call_kwargs = mock_client_instance.post.call_args + self.assertNotIn("keep_request_body", call_kwargs.get("params", {})) + + @patch("vespa.application.httpr.Client") + def test_request_body_none_via_session(self, MockClient): + self._mock_client(MockClient) + + app = Vespa(url="http://localhost", port=8080) + with app.syncio() as session: + r = session.query(body={"yql": "select * from sources * where true"}) + + self.assertIsNone(r.request_body) + + @patch("vespa.application.httpr.Client") + def test_request_body_stored_via_session(self, MockClient): + self._mock_client(MockClient) + body = {"yql": "select * from sources * where true"} + + app = Vespa(url="http://localhost", port=8080) + with app.syncio() as session: + r = session.query(body=body, keep_request_body=True) + + self.assertEqual(r.request_body, body) + self.assertIsNot(r.request_body, body) + + +@pytest.mark.asyncio +class TestAsyncQueryKeepRequestBody: + """Tests for the `keep_request_body` flag on `VespaAsync.query`.""" + + async def test_request_body_none_by_default(self): + app = Vespa(url="http://localhost", port=8080) + vespa_async = VespaAsync(app) + vespa_async._make_request = AsyncMock( + return_value=create_mock_httpr_response( + status_code=200, + json_data={"root": {}}, + url="http://localhost:8080/search/", + ) + ) + + r = await vespa_async.query(body={"yql": "select * from sources * where true"}) + + assert r.request_body is None + + async def test_request_body_stored_when_flag_true(self): + app = Vespa(url="http://localhost", port=8080) + vespa_async = VespaAsync(app) + vespa_async._make_request = AsyncMock( + return_value=create_mock_httpr_response( + status_code=200, + json_data={"root": {}}, + url="http://localhost:8080/search/", + ) + ) + body = {"yql": "select * from sources * where true", "nested": {"a": 1}} + + r = await vespa_async.query(body=body, keep_request_body=True) + + assert r.request_body == body + assert r.request_body is not body + assert r.request_body["nested"] is not body["nested"] + + # Mutating the original body after the call must not affect the stored copy. + body["nested"]["a"] = 999 + assert r.request_body["nested"]["a"] == 1 diff --git a/vespa/application.py b/vespa/application.py index 0cb5cec25..1edfc548f 100644 --- a/vespa/application.py +++ b/vespa/application.py @@ -2,6 +2,7 @@ import sys import asyncio +import copy import traceback import concurrent.futures import warnings @@ -573,6 +574,7 @@ def query( groupname: Optional[str] = None, streaming: bool = False, profile: bool = False, + keep_request_body: bool = False, **kwargs, ) -> Union[VespaQueryResponse, Generator[str, None, None]]: """ @@ -585,6 +587,8 @@ def query( groupname (str, optional): The groupname used with streaming search. streaming (bool, optional): Whether to use streaming mode (SSE). Defaults to False. profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False. + keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned + response as `response.request_body`. Defaults to False. Has no effect when streaming=True. **kwargs (dict, optional): Extra Vespa Query API parameters. Returns: @@ -605,6 +609,7 @@ def query( groupname=groupname, streaming=False, profile=profile, + keep_request_body=keep_request_body, **kwargs, ) @@ -1737,6 +1742,7 @@ def query( groupname: Optional[str] = None, streaming: bool = False, profile: bool = False, + keep_request_body: bool = False, **kwargs, ) -> Union[VespaQueryResponse, Generator[str, None, None]]: """ @@ -1747,6 +1753,8 @@ def query( groupname (str, optional): The groupname used in streaming search. streaming (bool, optional): Whether to use streaming mode (SSE). Defaults to False. profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False. + keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned + response as `response.request_body`. Defaults to False. Has no effect when streaming=True. **kwargs (dict, optional): Additional valid Vespa HTTP Query API parameters. See: . Returns: @@ -1776,6 +1784,7 @@ def query( json=response.json(), status_code=response.status_code, url=str(response.url), + request_body=copy.deepcopy(body) if keep_request_body else None, ) def _query_streaming( @@ -2325,6 +2334,7 @@ async def query( groupname: Optional[str] = None, profile: bool = False, retry_policy: Optional[AsyncRetrying] = None, + keep_request_body: bool = False, **kwargs, ) -> VespaQueryResponse: """ @@ -2335,6 +2345,8 @@ async def query( groupname (str, optional): The groupname used in streaming search. profile (bool, optional): Add profiling parameters to the query (response may be large). Defaults to False. retry_policy (AsyncRetrying, optional): Custom tenacity retry policy. Defaults to ``vespa.retries.QUERY_RETRY`` (five attempts with random exponential wait time). + keep_request_body (bool, optional): If True, store a deep copy of `body` on the returned + response as `response.request_body`. Defaults to False. **kwargs (dict, optional): Additional valid Vespa HTTP Query API parameters. Returns: @@ -2362,6 +2374,7 @@ async def _do_query() -> VespaQueryResponse: json=_response_json(response), status_code=response.status_code, url=str(response.url), + request_body=copy.deepcopy(body) if keep_request_body else None, ) return await retry_policy(_do_query) diff --git a/vespa/io.py b/vespa/io.py index 4bacd9283..bfbed0ef9 100644 --- a/vespa/io.py +++ b/vespa/io.py @@ -55,6 +55,8 @@ def __init__(self, json, status_code, url, request_body=None) -> None: @property def request_body(self) -> Optional[Dict]: + """A deep copy of the request body, set only when the query was made with + `keep_request_body=True`. `None` otherwise.""" return self._request_body @property From cb79b106ec3c2a7ccf98877b301fd15884d36206 Mon Sep 17 00:00:00 2001 From: Dainius Jocas Date: Sun, 20 Sep 2026 15:46:08 +0200 Subject: [PATCH 2/2] equality takes into account the request_body --- tests/unit/test_io.py | 65 ++++++++++++++++++++++++++++++++++++++++++- vespa/io.py | 11 +++++++- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py index 5fc2e77d0..6b03c67cc 100644 --- a/tests/unit/test_io.py +++ b/tests/unit/test_io.py @@ -1,7 +1,7 @@ # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. import unittest -from vespa.io import VespaVisitResponse, VespaQueryResponse +from vespa.io import VespaVisitResponse, VespaQueryResponse, VespaResponse class TestVespaVisitResult(unittest.TestCase): @@ -158,3 +158,66 @@ def test_hits(self): } ], ) + + +class TestVespaQueryResponseEquality(unittest.TestCase): + def setUp(self) -> None: + self.json = {"root": {"fields": {"totalCount": 1}}} + self.status_code = 200 + self.url = "http://localhost:8080/search/" + self.body = {"yql": "select * from sources * where true"} + + def _make(self, request_body=None): + return VespaQueryResponse( + json=self.json, + status_code=self.status_code, + url=self.url, + request_body=request_body, + ) + + def test_equal_when_request_body_matches(self): + self.assertEqual(self._make(), self._make()) + self.assertEqual( + self._make(request_body=self.body), self._make(request_body=self.body) + ) + + def test_equal_uses_value_equality_not_identity(self): + # Equal but distinct dict objects (e.g. one is a deep copy of the other). + a = self._make(request_body={"yql": "x", "nested": {"n": 1}}) + b = self._make(request_body={"yql": "x", "nested": {"n": 1}}) + self.assertIsNot(a.request_body, b.request_body) + self.assertEqual(a, b) + + def test_not_equal_when_request_body_differs(self): + a = self._make(request_body={"yql": "select * from sources * where a"}) + b = self._make(request_body={"yql": "select * from sources * where b"}) + self.assertNotEqual(a, b) + + def test_not_equal_when_one_request_body_is_none(self): + a = self._make(request_body=None) + b = self._make(request_body=self.body) + self.assertNotEqual(a, b) + self.assertNotEqual(b, a) + + def test_not_equal_to_unrelated_type(self): + a = self._make(request_body=self.body) + self.assertNotEqual(a, "not a response") + self.assertNotEqual( + a, + VespaResponse( + json=self.json, + status_code=self.status_code, + url=self.url, + operation_type="query", + ), + ) + + def test_base_fields_still_compared(self): + a = self._make(request_body=self.body) + b = VespaQueryResponse( + json={"root": {"fields": {"totalCount": 2}}}, + status_code=self.status_code, + url=self.url, + request_body=self.body, + ) + self.assertNotEqual(a, b) diff --git a/vespa/io.py b/vespa/io.py index bfbed0ef9..6602aea18 100644 --- a/vespa/io.py +++ b/vespa/io.py @@ -16,7 +16,10 @@ def __init__( self.operation_type = operation_type def __eq__(self, other: object) -> bool: - if not isinstance(other, self.__class__): + # Compare exact types (not isinstance) so equality stays symmetric + # across subclasses, e.g. a VespaQueryResponse is never equal to a + # plain VespaResponse even if their base fields happen to match. + if type(other) is not type(self): return NotImplemented return ( self.json == other.json @@ -53,6 +56,12 @@ def __init__(self, json, status_code, url, request_body=None) -> None: ) self._request_body = request_body + def __eq__(self, other: object) -> bool: + base_eq = super().__eq__(other) + if base_eq is NotImplemented: + return NotImplemented + return base_eq and self.request_body == other.request_body + @property def request_body(self) -> Optional[Dict]: """A deep copy of the request body, set only when the query was made with