Merge pull request #238 from yh-0/RDBC-670-1

RDBC-670 Streaming

Author: poissoncorp

ravendb/documents/commands/stream.py

@@ -0,0 +1,111 @@
+import requests
+from typing import TypeVar, Generic, Iterator
+
+from ravendb.documents.queries.index_query import IndexQuery
+from ravendb.documents.conventions import DocumentConventions
+from ravendb.extensions.json_extensions import JsonExtensions
+from ravendb.http.http_cache import HttpCache
+from ravendb.http.misc import ResponseDisposeHandling
+from ravendb.http.server_node import ServerNode
+from ravendb.http.raven_command import RavenCommand, RavenCommandResponseType
+from ravendb.json.metadata_as_dictionary import MetadataAsDictionary
+
+
+_T = TypeVar("_T")
+
+
+class StreamResult(Generic[_T]):
+    def __init__(
+        self, key: str = None, change_vector: str = None, metadata: MetadataAsDictionary = None, document: _T = None
+    ):
+        self.key = key
+        self.change_vector = change_vector
+        self.metadata = metadata
+        self.document = document
+
+
+class StreamResultResponse:
+    def __init__(self, response: requests.Response, stream_generator: Iterator):
+        self.response = response
+        self.stream_iterator = stream_generator
+
+
+class StreamCommand(RavenCommand[StreamResultResponse]):
+    def __init__(self, url: str):
+        super(StreamCommand, self).__init__(StreamResultResponse)
+        if not url or url.isspace():
+            raise ValueError("Url cannot be None or empty")
+
+        self._url = url
+        self._response_type = RavenCommandResponseType.EMPTY
+
+    def create_request(self, node: ServerNode) -> requests.Request:
+        return requests.Request("GET", f"{node.url}/databases/{node.database}/{self._url}")
+
+    def process_response(self, cache: HttpCache, response: requests.Response, url) -> ResponseDisposeHandling:
+        try:
+            result = StreamResultResponse(response, response.iter_lines())
+            self.result = result
+
+            return ResponseDisposeHandling.MANUALLY
+        except Exception as e:
+            raise RuntimeError("Unable to process stream response", e)
+
+    def send(self, session: requests.Session, request: requests.Request) -> requests.Response:
+        return session.request(
+            request.method,
+            url=request.url,
+            data=request.data,
+            files=request.files,
+            cert=session.cert,
+            headers=request.headers,
+            stream=True,
+        )
+
+    def is_read_request(self) -> bool:
+        return True
+
+
+class QueryStreamCommand(RavenCommand[StreamResultResponse]):
+    def __init__(self, conventions: DocumentConventions, query: IndexQuery):
+        super(QueryStreamCommand, self).__init__(StreamResultResponse)
+
+        if conventions is None:
+            raise ValueError("Conventions cannot be None")
+
+        if query is None:
+            raise ValueError("Query cannot be None")
+
+        self._conventions = conventions
+        self._index_query = query
+
+        self._response_type = RavenCommandResponseType.EMPTY
+
+    def create_request(self, node: ServerNode) -> requests.Request:
+        request = requests.Request("POST")
+        request.data = JsonExtensions.write_index_query(self._conventions, self._index_query)
+        request.url = f"{node.url}/databases/{node.database}/streams/queries?format=jsonl"
+        return request
+
+    def process_response(self, cache: HttpCache, response: requests.Response, url) -> ResponseDisposeHandling:
+        try:
+            stream_response = StreamResultResponse(response, response.iter_lines())
+            self.result = stream_response
+
+            return ResponseDisposeHandling.MANUALLY
+        except Exception as e:
+            raise RuntimeError("Unable to process stream response: " + e.args[0], e)
+
+    def send(self, session: requests.Session, request: requests.Request) -> requests.Response:
+        return session.request(
+            request.method,
+            url=request.url,
+            data=request.data,
+            files=request.files,
+            cert=session.cert,
+            headers=request.headers,
+            stream=True,
+        )
+
+    def is_read_request(self) -> bool:
+        return True

ravendb/documents/operations/executor.py

@@ -25,7 +25,7 @@
 
 
 _T = TypeVar("_T")
-_Operation_T = TypeVar("_Operation_T")
+_T_Operation_Return = TypeVar("_T_Operation_Return")
 
 
 class OperationExecutor:
@@ -41,7 +41,7 @@ def for_database(self, database_name: str) -> OperationExecutor:
             return self
         return OperationExecutor(self._store, database_name)
 
-    def send(self, operation: IOperation[_Operation_T], session_info: SessionInfo = None) -> _Operation_T:
+    def send(self, operation: IOperation[_T_Operation_Return], session_info: SessionInfo = None) -> _T_Operation_Return:
         command = operation.get_command(self._store, self._request_executor.conventions, self._request_executor.cache)
         self._request_executor.execute_command(command, session_info)
         return None if isinstance(operation, VoidOperation) else command.result
@@ -136,8 +136,8 @@ def for_database(self, database_name: str) -> MaintenanceOperationExecutor:
         return MaintenanceOperationExecutor(self._store, database_name)
 
     def send(
-        self, operation: Union[VoidMaintenanceOperation, MaintenanceOperation[_Operation_T]]
-    ) -> Optional[_Operation_T]:
+        self, operation: Union[VoidMaintenanceOperation, MaintenanceOperation[_T_Operation_Return]]
+    ) -> Optional[_T_Operation_Return]:
         self._assert_database_name_set()
         command = operation.get_command(self.request_executor.conventions)
         self.request_executor.execute_command(command)

ravendb/documents/session/document_session.py

@@ -8,7 +8,21 @@
 import os
 import time
 import uuid
-from typing import Union, Callable, TYPE_CHECKING, Optional, Dict, List, Type, TypeVar, Tuple, Generic, Set
+from typing import (
+    Union,
+    Callable,
+    TYPE_CHECKING,
+    Optional,
+    Dict,
+    List,
+    Type,
+    TypeVar,
+    Iterator,
+    Iterable,
+    Tuple,
+    Generic,
+    Set,
+)
 
 from ravendb.documents.session.document_session_revisions import DocumentSessionRevisions
 from ravendb.primitives import constants
@@ -22,7 +36,11 @@
     GetMultipleTimeSeriesOperation,
     TimeSeriesConfiguration,
 )
+from ravendb.documents.commands.stream import StreamResult
 from ravendb.documents.session.conditional_load import ConditionalLoadResult
+from ravendb.documents.session.operations.stream import StreamOperation
+from ravendb.documents.session.stream_statistics import StreamQueryStatistics
+from ravendb.documents.session.tokens.query_tokens.definitions import FieldsToFetchToken
 from ravendb.documents.session.time_series import (
     TimeSeriesEntry,
     TypedTimeSeriesEntry,
@@ -58,6 +76,7 @@
 from ravendb.documents.session.loaders.loaders import LoaderWithInclude, MultiLoaderWithInclude
 from ravendb.documents.session.operations.lazy import LazyLoadOperation, LazySessionOperations
 from ravendb.documents.session.operations.operations import MultiGetOperation, LoadStartingWithOperation
+from ravendb.documents.session.operations.query import QueryOperation
 from ravendb.documents.session.misc import (
     SessionOptions,
     ResponseTimeInformation,
@@ -67,7 +86,7 @@
     SessionInfo,
     TransactionMode,
 )
-from ravendb.documents.session.query import DocumentQuery, RawDocumentQuery
+from ravendb.documents.session.query import DocumentQuery, RawDocumentQuery, AbstractDocumentQuery
 from ravendb.json.metadata_as_dictionary import MetadataAsDictionary
 from ravendb.documents.session.operations.load_operation import LoadOperation
 from ravendb.tools.time_series import TSRangeHelper
@@ -1030,23 +1049,117 @@ def add_or_increment(self, key: str, entity: object, path_to_object: str, val_to
             patch_command_data.create_if_missing = new_instance
             self.defer(patch_command_data)
 
-        def stream(
-            self,
-            query_or_raw_query,  # : Union[RawDocumentQuery, DocumentQuery],
-            stream_query_stats,  # : Optional[StreamQueryStatistics] = None,
-        ) -> iter:
-            pass
+        def stream(self, query_or_raw_query: AbstractDocumentQuery[_T]) -> Iterator[StreamResult[_T]]:
+            stream_operation = StreamOperation(self._session)
+            command = stream_operation.create_request(query_or_raw_query.index_query)
+
+            self.request_executor.execute_command(command, self.session_info)
+
+            result = stream_operation.set_result(command.result)
+
+            return self._yield_result(query_or_raw_query, result)
+
+        def stream_with_statistics(
+            self, query_or_raw_query: AbstractDocumentQuery[_T], stats_callback: Callable[[StreamQueryStatistics], None]
+        ) -> Iterator[StreamResult[_T]]:
+            stats = StreamQueryStatistics()
+            stream_operation = StreamOperation(self._session, stats)
+            command = stream_operation.create_request(query_or_raw_query.index_query)
+
+            self.request_executor.execute_command(command, self.session_info)
+
+            result = stream_operation.set_result(command.result)
+
+            stats_callback(stats)
+
+            return self._yield_result(query_or_raw_query, result)
 
         def stream_starting_with(
             self,
-            object_type: type,
             starts_with: str,
-            matches: str = None,
+            matches: Optional[str] = None,
             start: int = 0,
-            page_size: int = 25,
-            starting_after: str = None,
-        ) -> iter:
-            pass
+            page_size: int = constants.int_max,
+            start_after: Optional[str] = None,
+            object_type: Optional[Type[_T]] = None,
+        ) -> Iterable[StreamResult[_T]]:
+            stream_operation = StreamOperation(self._session)
+
+            command = stream_operation.create_request_for_starts_with(
+                starts_with, matches, start, page_size, None, start_after
+            )
+            self.request_executor.execute_command(command, self.session_info)
+
+            result = stream_operation.set_result(command.result)
+            return DocumentSession._Advanced._StreamIterator(self, result, None, False, None, object_type)
+
+        class _StreamIterator(Iterator):
+            def __init__(
+                self,
+                session_advanced_reference: DocumentSession._Advanced,
+                inner_iterator: Iterator[Dict],
+                fields_to_fetch_token: Optional[FieldsToFetchToken],
+                is_project_into: bool,
+                on_next_item: Optional[Callable[[Dict], None]] = None,
+                object_type: Optional[Type[_T]] = None,
+            ):
+                self._session_advanced_reference = session_advanced_reference
+                self._inner_iterator = inner_iterator
+                self._fields_to_fetch_token = fields_to_fetch_token
+                self._is_project_into = is_project_into
+                self._on_next_item = on_next_item
+                self._object_type = object_type
+
+            def __iter__(self):
+                return self
+
+            def __next__(self):
+                next_value = self._inner_iterator.__next__()
+                try:
+                    if self._on_next_item is not None:
+                        self._on_next_item(next_value)
+                    return DocumentSession._Advanced._create_stream_result(
+                        self._session_advanced_reference,
+                        next_value,
+                        self._fields_to_fetch_token,
+                        self._is_project_into,
+                        self._object_type,
+                    )
+                except Exception as e:
+                    raise RuntimeError(f"Unable to parse stream result: {e.args[0]}", e)
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                pass
+
+        def _create_stream_result(
+            self,
+            json_dict: Dict,
+            fields_to_fetch: FieldsToFetchToken,
+            is_project_into: bool,
+            object_type: Optional[Type[_T]],
+        ) -> StreamResult[_T]:
+            metadata = json_dict.get(constants.Documents.Metadata.KEY)
+            change_vector = metadata.get(constants.Documents.Metadata.CHANGE_VECTOR)
+
+            # MapReduce indexes return reduce results tht don't have @id property
+            key = metadata.get(constants.Documents.Metadata.ID, None)
+
+            entity = QueryOperation.deserialize(
+                object_type, key, json_dict, metadata, fields_to_fetch, True, self._session, is_project_into
+            )
+
+            stream_result = StreamResult(key, change_vector, MetadataAsDictionary(metadata), entity)
+            return stream_result
+
+        def _yield_result(self, query: AbstractDocumentQuery, enumerator: Iterator[Dict]) -> Iterator[StreamResult[_T]]:
+            return DocumentSession._Advanced._StreamIterator(
+                self,
+                enumerator,
+                query._fields_to_fetch_token,
+                query.is_project_into,
+                query.invoke_after_stream_executed,
+                object_type=query.query_class,
+            )
 
         def stream_into(self):  # query: Union[DocumentQuery, RawDocumentQuery], output: iter):
             pass

ravendb/documents/session/operations/query.py

@@ -150,25 +150,25 @@ def deserialize(
         key: str,
         document: dict,
         metadata: dict,
-        fields_to_fetch: FieldsToFetchToken,
+        fields_to_fetch_token: FieldsToFetchToken,
         disable_entities_tracking: bool,
         session: "InMemoryDocumentSessionOperations",
         is_project_into: bool,
-    ):
+    ) -> _T:
         projection = metadata.get("@projection")
         if not projection:
             return session.track_entity(object_type, key, document, metadata, disable_entities_tracking)
 
         if (
-            fields_to_fetch is not None
-            and fields_to_fetch.projections is not None
-            and len(fields_to_fetch.projections) == 1
+            fields_to_fetch_token is not None
+            and fields_to_fetch_token.projections is not None
+            and len(fields_to_fetch_token.projections) == 1
         ):
-            projection_field = fields_to_fetch.projections[0]
+            projection_field = fields_to_fetch_token.projections[0]
 
-            if fields_to_fetch.source_alias is not None:
-                if projection_field.startswith(fields_to_fetch.source_alias):
-                    projection_field = projection_field[len(fields_to_fetch.source_alias) + 1 :]
+            if fields_to_fetch_token.source_alias is not None:
+                if projection_field.startswith(fields_to_fetch_token.source_alias):
+                    projection_field = projection_field[len(fields_to_fetch_token.source_alias) + 1 :]
 
                 if projection_field.startswith("'"):
                     projection_field = projection_field[1:-1]
@@ -188,8 +188,8 @@ def deserialize(
                     return Utils.get_default_value(object_type)
 
                 if (
-                    fields_to_fetch.fields_to_fetch is not None
-                    and fields_to_fetch.fields_to_fetch[0] == fields_to_fetch.projections[0]
+                    fields_to_fetch_token.fields_to_fetch is not None
+                    and fields_to_fetch_token.fields_to_fetch[0] == fields_to_fetch_token.projections[0]
                 ):
                     document = inner