RDBC-811 Make load_starting_with_into_stream work properly

Author: poissoncorp

ravendb/documents/commands/results.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import Union, Optional, List, Dict
+from typing import Union, Optional, List, Dict, Any
 
 
 class GetDocumentResult:
@@ -39,3 +39,13 @@ def from_json(cls, json_dict: dict) -> GetDocumentsResult:
             json_dict.get("CompareExchangeValueIncludes", None),
             json_dict.get("NextPageStart", None),
         )
+
+    def to_json(self) -> Dict[str, Any]:
+        return {
+            "Includes": self.includes,
+            "Results": self.results,
+            "CounterIncludes": self.counter_includes,
+            "TimeSeriesIncludes": self.time_series_includes,
+            "CompareExchangeValueIncludes": self.compare_exchange_includes,
+            "NextPageStart": self.next_page_start,
+        }

ravendb/documents/session/document_session.py

@@ -375,16 +375,30 @@ def load_starting_with(
         start_after: Optional[str] = None,
     ) -> List[_T]:
         load_starting_with_operation = LoadStartingWithOperation(self)
-        self.__load_starting_with_internal(
-            id_prefix, load_starting_with_operation, None, matches, start, page_size, exclude, start_after
+        self._load_starting_with_internal(
+            id_prefix, load_starting_with_operation, matches, start, page_size, exclude, start_after
         )
         return load_starting_with_operation.get_documents(object_type)
 
-    def __load_starting_with_internal(
+    def load_starting_with_into_stream(
+        self,
+        id_prefix: str,
+        matches: str = None,
+        start: int = 0,
+        page_size: int = 25,
+        exclude: str = None,
+        start_after: str = None,
+    ) -> bytes:
+        if id_prefix is None:
+            raise ValueError("Arg 'id_prefix' is cannot be None.")
+        return self._load_starting_with_into_stream_internal(
+            id_prefix, LoadStartingWithOperation(self), matches, start, page_size, exclude, start_after
+        )
+
+    def _load_starting_with_internal(
         self,
         id_prefix: str,
         operation: LoadStartingWithOperation,
-        stream,
         matches: str,
         start: int,
         page_size: int,
@@ -395,12 +409,31 @@ def __load_starting_with_internal(
         command = operation.create_request()
         if command:
             self._request_executor.execute_command(command, self.session_info)
-            if stream:
-                pass  # todo: stream
-            else:
-                operation.set_result(command.result)
+            operation.set_result(command.result)
         return command
 
+    def _load_starting_with_into_stream_internal(
+        self,
+        id_prefix: str,
+        operation: LoadStartingWithOperation,
+        matches: str,
+        start: int,
+        page_size: int,
+        exclude: str,
+        start_after: str,
+    ) -> bytes:
+        operation.with_start_with(id_prefix, matches, start, page_size, exclude, start_after)
+        command = operation.create_request()
+        bytes_result = None
+        if command:
+            self.request_executor.execute_command(command, self.session_info)
+            try:
+                result = command.result
+                bytes_result = json.dumps(result.to_json()).encode("utf-8")
+            except Exception as e:
+                raise RuntimeError("Unable sto serialize returned value into stream") from e
+        return bytes_result
+
     def document_query_from_index_type(self, index_type: Type[_TIndex], object_type: Type[_T]) -> DocumentQuery[_T]:
         try:
             index = Utils.try_get_new_instance(index_type)
@@ -783,32 +816,6 @@ def wait_for_indexes_after_save_changes(
 
             index_options.wait_for_indexes = True
 
-        def __load_starting_with_internal(
-            self,
-            id_prefix: str,
-            operation: LoadStartingWithOperation,
-            stream: Union[None, bytes],
-            matches: str,
-            start: int,
-            page_size: int,
-            exclude: str,
-            start_after: str,
-        ) -> GetDocumentsCommand:
-            operation.with_start_with(id_prefix, matches, start, page_size, exclude, start_after)
-            command = operation.create_request()
-            if command is not None:
-                self._session._request_executor.execute(command, self._session.session_info)
-                if stream:
-                    try:
-                        result = command.result
-                        stream_to_dict = json.loads(stream.decode("utf-8"))
-                        result.__dict__.update(stream_to_dict)
-                    except IOError as e:
-                        raise RuntimeError(f"Unable to serialize returned value into stream {e.args[0]}", e)
-                else:
-                    operation.set_result(command.result)
-            return command
-
         def load_starting_with(
             self,
             id_prefix: str,
@@ -826,26 +833,14 @@ def load_starting_with(
         def load_starting_with_into_stream(
             self,
             id_prefix: str,
-            output: bytes,
             matches: str = None,
             start: int = 0,
             page_size: int = 25,
             exclude: str = None,
             start_after: str = None,
-        ):
-            if not output:
-                raise ValueError("Output cannot be None")
-            if not id_prefix:
-                raise ValueError("Id prefix cannot be None")
-            self.__load_starting_with_internal(
-                id_prefix,
-                LoadStartingWithOperation(self._session),
-                output,
-                matches,
-                start,
-                page_size,
-                exclude,
-                start_after,
+        ) -> bytes:
+            return self._session.load_starting_with_into_stream(
+                id_prefix, matches, start, page_size, exclude, start_after
             )
 
         def load_into_stream(self, keys: List[str], output: bytes) -> None:

ravendb/tests/jvm_migrated_tests/client_tests/test_load_into_stream.py

@@ -0,0 +1,40 @@
+import json
+
+from ravendb import DocumentStore
+from ravendb.documents.commands.results import GetDocumentsResult
+from ravendb.infrastructure.orders import Employee
+from ravendb.tests.test_base import TestBase
+
+
+class TestLoadIntoStream(TestBase):
+    def setUp(self):
+        super().setUp()
+
+    @staticmethod
+    def insert_data(store: DocumentStore):
+        with store.open_session() as session:
+
+            def _insert_employee(name: str = None):
+                employee = Employee(first_name=name)
+                session.store(employee)
+
+            _insert_employee("Aviv")
+            _insert_employee("Iftah")
+            _insert_employee("Tal")
+            _insert_employee("Maxim")
+            _insert_employee("Karmel")
+            _insert_employee("Grisha")
+            _insert_employee("Michael")
+            session.save_changes()
+
+    def test_can_load_starting_with_into_stream(self):
+        self.insert_data(self.store)
+        with self.store.open_session() as session:
+            stream = session.advanced.load_starting_with_into_stream("employees/")
+            json_node = json.loads(stream.decode("utf-8"))
+            result = GetDocumentsResult.from_json(json_node)
+            self.assertEqual(7, len(result.results))
+            names = ["Aviv", "Iftah", "Tal", "Maxim", "Karmel", "Grisha", "Michael"]
+            for name_from_results in [result["first_name"] for result in result.results]:
+                self.assertIn(name_from_results, names)
+                names.remove(name_from_results)