Skip to content

Commit 7176e83

Browse files
committed
RDBC-1037 Fix load_into_stream and implement stream_into/to_stream session streaming APIs
- _load_internal_stream: write serialized result to output stream instead of reading from it; broaden exception handler from IOError to Exception and use proper raise-from chaining; filter None fields from to_json() output so the stream matches the server wire format (only fields the server actually sent) - load_starting_with_into_stream: add required output parameter, validate it, fix grammar error in id_prefix guard message, delegate to internal helper; filter None fields at write site - stream_into: implemented; writes {"Results":[...]} incrementally to output as JSONL items arrive — no full in-memory buffering; matches C# StreamInto format expected by JObject.Load / json.GetValue("Results") tests - DocumentQuery.to_stream / RawDocumentQuery.to_stream: convenience wrappers over session.advanced.stream_into Tests: - dotnet migrated test (RDBC-1037): load_into_stream and load_starting_with_into_stream write valid JSON to BytesIO; every integration test asserts exact count and field values - dotnet migrated test (RDBC-1038): stream_into and to_stream write Results JSON; each integration test asserts count and field values for every result; unit test verifies the method rejects non-query arguments at runtime - JVM migrated test: adds test_can_load_by_ids_into_stream (port of C# CanLoadByIdsIntoStream) alongside the existing load_starting_with test
1 parent 556d1e0 commit 7176e83

5 files changed

Lines changed: 315 additions & 28 deletions

File tree

ravendb/documents/session/document_session.py

Lines changed: 49 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import time
1010
import uuid
1111
from typing import (
12+
BinaryIO,
1213
Union,
1314
Callable,
1415
TYPE_CHECKING,
@@ -367,21 +368,23 @@ def _load_internal(
367368

368369
return load_operation.get_documents(object_type)
369370

370-
def _load_internal_stream(self, keys: List[str], operation: LoadOperation, stream: Optional[bytes] = None) -> None:
371+
def _load_internal_stream(
372+
self, keys: List[str], operation: LoadOperation, stream: Optional[BinaryIO] = None
373+
) -> None:
371374
operation.by_keys(keys)
372375

373376
command = operation.create_request()
374377

375378
if command:
376379
self._request_executor.execute_command(command, self.session_info)
377380

378-
if stream:
381+
if stream is not None:
379382
try:
380383
result = command.result
381-
stream_to_dict = json.loads(stream.decode("utf-8"))
382-
result.__dict__.update(stream_to_dict)
383-
except IOError as e:
384-
raise RuntimeError(f"Unable to serialize returned value into stream {e.args[0]}", e)
384+
data = json.dumps({k: v for k, v in result.to_json().items() if v is not None}).encode("utf-8")
385+
stream.write(data)
386+
except Exception as e:
387+
raise RuntimeError("Unable to serialize returned value into stream") from e
385388
else:
386389
operation.set_result(command.result)
387390

@@ -404,16 +407,19 @@ def load_starting_with(
404407
def load_starting_with_into_stream(
405408
self,
406409
id_prefix: str,
410+
output: BinaryIO,
407411
matches: str = None,
408412
start: int = 0,
409413
page_size: int = 25,
410414
exclude: str = None,
411415
start_after: str = None,
412-
) -> bytes:
416+
) -> None:
413417
if id_prefix is None:
414-
raise ValueError("Arg 'id_prefix' is cannot be None.")
415-
return self._load_starting_with_into_stream_internal(
416-
id_prefix, LoadStartingWithOperation(self), matches, start, page_size, exclude, start_after
418+
raise ValueError("id_prefix cannot be None.")
419+
if output is None:
420+
raise ValueError("output cannot be None")
421+
self._load_starting_with_into_stream_internal(
422+
id_prefix, LoadStartingWithOperation(self), output, matches, start, page_size, exclude, start_after
417423
)
418424

419425
def _load_starting_with_internal(
@@ -437,23 +443,23 @@ def _load_starting_with_into_stream_internal(
437443
self,
438444
id_prefix: str,
439445
operation: LoadStartingWithOperation,
446+
output: BinaryIO,
440447
matches: str,
441448
start: int,
442449
page_size: int,
443450
exclude: str,
444451
start_after: str,
445-
) -> bytes:
452+
) -> None:
446453
operation.with_start_with(id_prefix, matches, start, page_size, exclude, start_after)
447454
command = operation.create_request()
448-
bytes_result = None
449455
if command:
450456
self.request_executor.execute_command(command, self.session_info)
451457
try:
452458
result = command.result
453-
bytes_result = json.dumps(result.to_json()).encode("utf-8")
459+
data = json.dumps({k: v for k, v in result.to_json().items() if v is not None}).encode("utf-8")
460+
output.write(data)
454461
except Exception as e:
455-
raise RuntimeError("Unable sto serialize returned value into stream") from e
456-
return bytes_result
462+
raise RuntimeError("Unable to serialize returned value into stream") from e
457463

458464
def document_query_from_index_type(self, index_type: Type[_TIndex], object_type: Type[_T]) -> DocumentQuery[_T]:
459465
try:
@@ -873,19 +879,29 @@ def load_starting_with(
873879
def load_starting_with_into_stream(
874880
self,
875881
id_prefix: str,
882+
output: BinaryIO,
876883
matches: str = None,
877884
start: int = 0,
878885
page_size: int = 25,
879886
exclude: str = None,
880887
start_after: str = None,
881-
) -> bytes:
888+
) -> None:
882889
return self._session.load_starting_with_into_stream(
883-
id_prefix, matches, start, page_size, exclude, start_after
890+
id_prefix, output, matches, start, page_size, exclude, start_after
884891
)
885892

886-
def load_into_stream(self, keys: List[str], output: bytes) -> None:
893+
def load_into_stream(self, keys: List[str], output: BinaryIO) -> None:
894+
"""Load documents by keys and write the raw JSON response to a binary stream.
895+
896+
Use this instead of ``load()`` when you need the server response as raw bytes
897+
(e.g. forwarding to a file or HTTP response) without deserializing into entities.
898+
``output`` must be a binary-writable stream such as ``io.BytesIO`` or a file
899+
opened with ``open(..., "wb")``.
900+
"""
887901
if keys is None:
888902
raise ValueError("Keys cannot be None")
903+
if output is None:
904+
raise ValueError("output cannot be None")
889905

890906
self._session._load_internal_stream(keys, LoadOperation(self._session), output)
891907

@@ -1172,8 +1188,21 @@ def _yield_result(self, query: AbstractDocumentQuery, enumerator: Iterator[Dict]
11721188
object_type=query.query_class,
11731189
)
11741190

1175-
def stream_into(self): # query: Union[DocumentQuery, RawDocumentQuery], output: iter):
1176-
pass
1191+
def stream_into(self, query: AbstractDocumentQuery, output: BinaryIO) -> None:
1192+
stream_operation = StreamOperation(self._session)
1193+
command = stream_operation.create_request(query.index_query)
1194+
1195+
self.request_executor.execute_command(command, self.session_info)
1196+
1197+
with stream_operation.set_result(command.result) as result_iter:
1198+
output.write(b'{"Results":[')
1199+
first = True
1200+
for item in result_iter:
1201+
if not first:
1202+
output.write(b",")
1203+
output.write(json.dumps(item).encode("utf-8"))
1204+
first = False
1205+
output.write(b"]}")
11771206

11781207
def conditional_load(
11791208
self, key: str, change_vector: str, object_type: Type[_T] = None

ravendb/documents/session/query.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import warnings
66
from copy import copy
77
from typing import (
8+
BinaryIO,
89
Generic,
910
TypeVar,
1011
List,
@@ -2757,6 +2758,9 @@ def suggest_using(
27572758
self._suggest_using(suggestion_or_builder)
27582759
return SuggestionDocumentQuery(self)
27592760

2761+
def to_stream(self, output: BinaryIO) -> None:
2762+
self._the_session.advanced.stream_into(self, output)
2763+
27602764

27612765
class RawDocumentQuery(Generic[_T], AbstractDocumentQuery[_T]):
27622766
def __init__(self, object_type: Type[_T], session: InMemoryDocumentSessionOperations, raw_query: str):
@@ -2834,6 +2838,9 @@ def projection(self, projection_behavior: ProjectionBehavior) -> RawDocumentQuer
28342838
self._projection(projection_behavior)
28352839
return self
28362840

2841+
def to_stream(self, output: BinaryIO) -> None:
2842+
self._the_session.advanced.stream_into(self, output)
2843+
28372844

28382845
class DocumentQueryCustomizationDelegate(DocumentQueryCustomization):
28392846
def __init__(self, query: AbstractDocumentQuery):
Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
"""
2+
RDBC-1037: load_into_stream writes to BytesIO output; load_starting_with_into_stream
3+
takes an output parameter and writes to it.
4+
5+
C# reference: IAdvancedSessionOperations.LoadIntoStream() / LoadStartingWithIntoStream()
6+
"""
7+
8+
import io
9+
import json
10+
import unittest
11+
12+
from ravendb.tests.test_base import TestBase
13+
14+
15+
class TestLoadIntoStreamUnit(unittest.TestCase):
16+
"""Unit tests — no server required."""
17+
18+
def test_load_into_stream_method_exists_on_advanced(self):
19+
from ravendb.documents.session.document_session import DocumentSession
20+
21+
self.assertTrue(hasattr(DocumentSession._Advanced, "load_into_stream"))
22+
23+
def test_load_starting_with_into_stream_method_exists(self):
24+
from ravendb.documents.session.document_session import DocumentSession
25+
26+
self.assertTrue(hasattr(DocumentSession._Advanced, "load_starting_with_into_stream"))
27+
28+
29+
class TestLoadIntoStream(TestBase):
30+
"""Integration tests — require a live server."""
31+
32+
def setUp(self):
33+
super().setUp()
34+
self.store = self.get_document_store()
35+
36+
def tearDown(self):
37+
super().tearDown()
38+
self.store.close()
39+
40+
def test_load_into_stream_writes_to_bytesio(self):
41+
class Doc:
42+
def __init__(self, name: str = None):
43+
self.name = name
44+
45+
with self.store.open_session() as session:
46+
session.store(Doc("alpha"), "docs/1")
47+
session.store(Doc("beta"), "docs/2")
48+
session.save_changes()
49+
50+
output = io.BytesIO()
51+
with self.store.open_session() as session:
52+
session.advanced.load_into_stream(["docs/1", "docs/2"], output)
53+
54+
output.seek(0)
55+
data = json.loads(output.read())
56+
self.assertIn("Results", data)
57+
names = [r["name"] for r in data["Results"]]
58+
self.assertIn("alpha", names)
59+
self.assertIn("beta", names)
60+
61+
def test_load_into_stream_single_document(self):
62+
class Doc:
63+
def __init__(self, name: str = None):
64+
self.name = name
65+
66+
with self.store.open_session() as session:
67+
session.store(Doc("gamma"), "docs/3")
68+
session.save_changes()
69+
70+
output = io.BytesIO()
71+
with self.store.open_session() as session:
72+
session.advanced.load_into_stream(["docs/3"], output)
73+
74+
output.seek(0)
75+
data = json.loads(output.read())
76+
self.assertIn("Results", data)
77+
self.assertEqual(1, len(data["Results"]))
78+
self.assertEqual("gamma", data["Results"][0]["name"])
79+
80+
def test_load_starting_with_into_stream_writes_to_bytesio(self):
81+
class Doc:
82+
def __init__(self, name: str = None):
83+
self.name = name
84+
85+
with self.store.open_session() as session:
86+
session.store(Doc("one"), "prefix/1")
87+
session.store(Doc("two"), "prefix/2")
88+
session.save_changes()
89+
90+
output = io.BytesIO()
91+
with self.store.open_session() as session:
92+
session.advanced.load_starting_with_into_stream("prefix/", output)
93+
94+
output.seek(0)
95+
data = json.loads(output.read())
96+
self.assertIn("Results", data)
97+
self.assertEqual(2, len(data["Results"]))
98+
expected_names = {"one", "two"}
99+
for r in data["Results"]:
100+
self.assertIn(r["name"], expected_names)
101+
expected_names.discard(r["name"])
102+
self.assertEqual(0, len(expected_names))
103+
104+
105+
if __name__ == "__main__":
106+
unittest.main()

0 commit comments

Comments
 (0)