PYTHON-1362 - Add find/aggregate_raw_batches()

Rename find_raw to find_raw_batches, and add aggregate_raw_batches.
Rename RawBSONCursor and RawBSONCommandCursor to RawBatchCursor and
RawBatchCommandCursor.

Author: ajdavis

Diff for: doc/api/pymongo/cursor.rst

@@ -24,4 +24,4 @@
 
       .. automethod:: __getitem__
 
-   .. autoclass:: pymongo.cursor.RawBSONCursor(collection, filter=None, projection=None, skip=0, limit=0, no_cursor_timeout=False, cursor_type=CursorType.NON_TAILABLE, sort=None, allow_partial_results=False, oplog_replay=False, modifiers=None, batch_size=0, collation=None, hint=None, max_scan=None, max_time_ms=None, max=None, min=None, return_key=False, show_record_id=False, snapshot=False, comment=None)
+   .. autoclass:: pymongo.cursor.RawBatchCursor(collection, filter=None, projection=None, skip=0, limit=0, no_cursor_timeout=False, cursor_type=CursorType.NON_TAILABLE, sort=None, allow_partial_results=False, oplog_replay=False, modifiers=None, batch_size=0, collation=None, hint=None, max_scan=None, max_time_ms=None, max=None, min=None, return_key=False, show_record_id=False, snapshot=False, comment=None)

Diff for: doc/changelog.rst

@@ -8,6 +8,12 @@ This version drops support for MongoDB versions older than 2.6. If connecting to
 a MongoDB 2.4 server or older, PyMongo now throws a
 :exc:`~pymongo.errors.ConfigurationError`.
 
+Highlights include:
+
+- New methods :meth:`~pymongo.collection.Collection.find_raw_batches` and
+  :meth:`~pymongo.collection.Collection.aggregate_raw_batches` for use with
+  external libraries that can parse raw batches of BSON data.
+
 Changes in Version 3.5.1
 ------------------------
 

Diff for: pymongo/collection.py

@@ -30,9 +30,9 @@
                      helpers,
                      message)
 from pymongo.bulk import BulkOperationBuilder, _Bulk
-from pymongo.command_cursor import CommandCursor
+from pymongo.command_cursor import CommandCursor, RawBatchCommandCursor
 from pymongo.collation import validate_collation_or_none
-from pymongo.cursor import Cursor, RawBSONCursor
+from pymongo.cursor import Cursor, RawBatchCursor
 from pymongo.errors import ConfigurationError, InvalidName, OperationFailure
 from pymongo.helpers import _check_write_command_response
 from pymongo.helpers import _UNICODE_REPLACE_CODEC_OPTIONS
@@ -1278,25 +1278,25 @@ def find(self, *args, **kwargs):
         """
         return Cursor(self, *args, **kwargs)
 
-    def find_raw(self, *args, **kwargs):
+    def find_raw_batches(self, *args, **kwargs):
         """Query the database and retrieve batches of raw BSON.
 
         Takes the same parameters as :meth:`find` but returns a
-        :class:`~pymongo.cursor.RawBSONCursor`.
+        :class:`~pymongo.cursor.RawBatchCursor`.
 
         This example demonstrates how to work with raw batches, but in practice
         raw batches should be passed to an external library that can decode
         BSON into another data type, rather than used with PyMongo's
         :mod:`bson` module.
 
           >>> import bson
-          >>> cursor = db.test.find_raw()
+          >>> cursor = db.test.find_raw_batches()
           >>> for batch in cursor:
           ...     print(bson.decode_all(batch))
 
         .. versionadded:: 3.6
         """
-        return RawBSONCursor(self, *args, **kwargs)
+        return RawBatchCursor(self, *args, **kwargs)
 
     def parallel_scan(self, num_cursors, **kwargs):
         """Scan this entire collection in parallel.
@@ -1821,6 +1821,69 @@ def options(self):
 
         return options
 
+    def _aggregate(self, pipeline, cursor_class, first_batch_size, **kwargs):
+        if not isinstance(pipeline, list):
+            raise TypeError("pipeline must be a list")
+
+        if "explain" in kwargs:
+            raise ConfigurationError("The explain option is not supported. "
+                                     "Use Database.command instead.")
+        collation = validate_collation_or_none(kwargs.pop('collation', None))
+
+        cmd = SON([("aggregate", self.__name),
+                   ("pipeline", pipeline)])
+
+        # Remove things that are not command options.
+        batch_size = common.validate_non_negative_integer_or_none(
+            "batchSize", kwargs.pop("batchSize", None))
+        use_cursor = common.validate_boolean(
+            "useCursor", kwargs.pop("useCursor", True))
+        # If the server does not support the "cursor" option we
+        # ignore useCursor and batchSize.
+        with self._socket_for_reads() as (sock_info, slave_ok):
+            if sock_info.max_wire_version > 0:
+                if use_cursor:
+                    if "cursor" not in kwargs:
+                        kwargs["cursor"] = {}
+                    if first_batch_size is not None:
+                        kwargs["cursor"]["batchSize"] = first_batch_size
+
+            dollar_out = pipeline and '$out' in pipeline[-1]
+            if (sock_info.max_wire_version >= 5 and dollar_out and
+                    self.write_concern):
+                cmd['writeConcern'] = self.write_concern.document
+
+            cmd.update(kwargs)
+
+            # Apply this Collection's read concern if $out is not in the
+            # pipeline.
+            if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
+                if dollar_out:
+                    result = self._command(sock_info, cmd, slave_ok,
+                                           parse_write_concern_error=True,
+                                           collation=collation)
+                else:
+                    result = self._command(sock_info, cmd, slave_ok,
+                                           read_concern=self.read_concern,
+                                           collation=collation)
+            else:
+                result = self._command(sock_info, cmd, slave_ok,
+                                       parse_write_concern_error=dollar_out,
+                                       collation=collation)
+
+            if "cursor" in result:
+                cursor = result["cursor"]
+            else:
+                # Pre-MongoDB 2.6. Fake a cursor.
+                cursor = {
+                    "id": 0,
+                    "firstBatch": result["result"],
+                    "ns": self.full_name,
+                }
+
+            return cursor_class(
+                self, cursor, sock_info.address).batch_size(batch_size or 0)
+
     def aggregate(self, pipeline, **kwargs):
         """Perform an aggregation using the aggregation framework on this
         collection.
@@ -1892,66 +1955,34 @@ def aggregate(self, pipeline, **kwargs):
         .. _aggregate command:
             http://docs.mongodb.org/manual/applications/aggregation
         """
-        if not isinstance(pipeline, list):
-            raise TypeError("pipeline must be a list")
-
-        if "explain" in kwargs:
-            raise ConfigurationError("The explain option is not supported. "
-                                     "Use Database.command instead.")
-        collation = validate_collation_or_none(kwargs.pop('collation', None))
+        return self._aggregate(pipeline,
+                               CommandCursor,
+                               kwargs.get('batchSize'),
+                               **kwargs)
 
-        cmd = SON([("aggregate", self.__name),
-                   ("pipeline", pipeline)])
-
-        # Remove things that are not command options.
-        batch_size = common.validate_positive_integer_or_none(
-            "batchSize", kwargs.pop("batchSize", None))
-        use_cursor = common.validate_boolean(
-            "useCursor", kwargs.pop("useCursor", True))
-        # If the server does not support the "cursor" option we
-        # ignore useCursor and batchSize.
-        with self._socket_for_reads() as (sock_info, slave_ok):
-            if sock_info.max_wire_version > 0:
-                if use_cursor:
-                    if "cursor" not in kwargs:
-                        kwargs["cursor"] = {}
-                    if batch_size is not None:
-                        kwargs["cursor"]["batchSize"] = batch_size
+    def aggregate_raw_batches(self, pipeline, **kwargs):
+        """Perform an aggregation and retrieve batches of raw BSON.
 
-            dollar_out = pipeline and '$out' in pipeline[-1]
-            if (sock_info.max_wire_version >= 5 and dollar_out and
-                    self.write_concern):
-                cmd['writeConcern'] = self.write_concern.document
+        Takes the same parameters as :meth:`aggregate` but returns a
+        :class:`~pymongo.cursor.RawBatchCursor`.
 
-            cmd.update(kwargs)
+        This example demonstrates how to work with raw batches, but in practice
+        raw batches should be passed to an external library that can decode
+        BSON into another data type, rather than used with PyMongo's
+        :mod:`bson` module.
 
-            # Apply this Collection's read concern if $out is not in the
-            # pipeline.
-            if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
-                if dollar_out:
-                    result = self._command(sock_info, cmd, slave_ok,
-                                           parse_write_concern_error=True,
-                                           collation=collation)
-                else:
-                    result = self._command(sock_info, cmd, slave_ok,
-                                           read_concern=self.read_concern,
-                                           collation=collation)
-            else:
-                result = self._command(sock_info, cmd, slave_ok,
-                                       parse_write_concern_error=dollar_out,
-                                       collation=collation)
+          >>> import bson
+          >>> cursor = db.test.aggregate_raw_batches([
+          ...     {'$project': {'x': {'$multiply': [2, '$x']}}}])
+          >>> for batch in cursor:
+          ...     print(bson.decode_all(batch))
 
-            if "cursor" in result:
-                cursor = result["cursor"]
-            else:
-                # Pre-MongoDB 2.6. Fake a cursor.
-                cursor = {
-                    "id": 0,
-                    "firstBatch": result["result"],
-                    "ns": self.full_name,
-                }
-            return CommandCursor(
-                self, cursor, sock_info.address).batch_size(batch_size or 0)
+        .. versionadded:: 3.6
+        """
+        return self._aggregate(pipeline,
+                               RawBatchCommandCursor,
+                               0,
+                               **kwargs)
 
     def group(self, key, condition, initial, reduce, finalize=None, **kwargs):
         """Perform a query similar to an SQL *group by* operation.

Diff for: pymongo/command_cursor.py

@@ -20,22 +20,29 @@
 
 from bson.py3compat import integer_types
 from pymongo import helpers
-from pymongo.errors import AutoReconnect, NotMasterError, OperationFailure
-from pymongo.message import _CursorAddress, _GetMore, _convert_exception
+from pymongo.errors import (AutoReconnect,
+                            InvalidOperation,
+                            NotMasterError,
+                            OperationFailure)
+from pymongo.message import (_convert_exception,
+                             _CursorAddress,
+                             _GetMore,
+                             _RawBatchGetMore)
 
 
 class CommandCursor(object):
-    """A cursor / iterator over command cursors.
-    """
+    """A cursor / iterator over command cursors."""
+    _getmore_class = _GetMore
 
     def __init__(self, collection, cursor_info, address, retrieved=0):
         """Create a new command cursor.
+
+        The parameter 'retrieved' is unused.
         """
         self.__collection = collection
         self.__id = cursor_info['id']
         self.__address = address
         self.__data = deque(cursor_info['firstBatch'])
-        self.__retrieved = retrieved
         self.__batch_size = 0
         self.__killed = (self.__id == 0)
 
@@ -115,9 +122,9 @@ def __send_message(self, operation):
         if publish:
             start = datetime.datetime.now()
         try:
-            doc = helpers._unpack_response(response.data,
-                                           self.__id,
-                                           self.__collection.codec_options)
+            doc = self._unpack_response(response.data,
+                                        self.__id,
+                                        self.__collection.codec_options)
             if from_command:
                 helpers._check_command_response(doc['data'][0])
 
@@ -154,11 +161,9 @@ def __send_message(self, operation):
             cursor = doc['data'][0]['cursor']
             documents = cursor['nextBatch']
             self.__id = cursor['id']
-            self.__retrieved += len(documents)
         else:
             documents = doc["data"]
             self.__id = doc["cursor_id"]
-            self.__retrieved += doc["number_returned"]
 
         if publish:
             duration = (datetime.datetime.now() - start) + cmd_duration
@@ -174,6 +179,8 @@ def __send_message(self, operation):
             self.__killed = True
         self.__data = deque(documents)
 
+    def _unpack_response(self, response, cursor_id, codec_options):
+        return helpers._unpack_response(response, cursor_id, codec_options)
 
     def _refresh(self):
         """Refreshes the cursor with more data from the server.
@@ -188,17 +195,21 @@ def _refresh(self):
         if self.__id:  # Get More
             dbname, collname = self.__ns.split('.', 1)
             self.__send_message(
-                _GetMore(dbname,
-                         collname,
-                         self.__batch_size,
-                         self.__id,
-                         self.__collection.codec_options))
+                self._getmore_class(dbname,
+                                    collname,
+                                    self.__batch_size,
+                                    self.__id,
+                                    self.__collection.codec_options))
 
         else:  # Cursor id is zero nothing else to return
             self.__killed = True
 
         return len(self.__data)
 
+    @property
+    def _collection(self):
+        return self.__collection
+
     @property
     def alive(self):
         """Does this cursor have the potential to return more data?
@@ -247,3 +258,31 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
+
+
+class RawBatchCommandCursor(CommandCursor):
+    _getmore_class = _RawBatchGetMore
+
+    def __init__(self, collection, cursor_info, address, retrieved=0):
+        """Create a new cursor / iterator over raw batches of BSON data.
+
+        Should not be called directly by application developers -
+        see :meth:`~pymongo.collection.Collection.aggregate_raw_batches`
+        instead.
+
+        .. mongodoc:: cursors
+        """
+        assert not cursor_info.get('firstBatch')
+        super(RawBatchCommandCursor, self).__init__(
+            collection, cursor_info, address, retrieved)
+
+        db = self._collection.database
+        if db.outgoing_manipulators or db.outgoing_copying_manipulators:
+            raise InvalidOperation("Raw batches are not compatible with"
+                                   " SON manipulators.")
+
+    def _unpack_response(self, response, cursor_id, codec_options):
+        return helpers._raw_response(response, cursor_id)
+
+    def __getitem__(self, index):
+        raise InvalidOperation("Cannot call __getitem__ on RawBatchCursor")