PYTHON-1866 Add support for $merge aggregation pipeline stage

prashantmital · prashantmital · commit 694a4a5d853f · 2019-06-13T19:38:15.000-07:00
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -49,6 +49,8 @@ Version 3.9 adds support for MongoDB 4.2. Highlights include:
 - Support for publishing Connection Monitoring and Pooling events via the new
   :class:`~pymongo.monitoring.ConnectionPoolListener` class. See
   :mod:`~pymongo.monitoring` for an example.
+- :meth:`pymongo.collection.Collection.aggregate` and
+  :meth:`pymongo.database.Database.aggregate` now support the ``$merge`` pipeline
 
 .. _URI options specification: https://github.com/mongodb/specifications/blob/master/source/uri-options/uri-options.rst
 
diff --git a/pymongo/aggregation.py b/pymongo/aggregation.py
@@ -38,10 +38,25 @@ def __init__(self, target, cursor_class, pipeline, options,
 
         common.validate_list('pipeline', pipeline)
         self._pipeline = pipeline
+        self._performs_write = False
+        if pipeline and ("$out" in pipeline[-1] or "$merge" in pipeline[-1]):
+            self._performs_write = True
 
         common.validate_is_mapping('options', options)
         self._options = options
 
+        # This is the batchSize that will be used for setting the initial
+        # batchSize for the cursor, as well as the subsequent getMores.
+        self._batch_size = common.validate_non_negative_integer_or_none(
+            "batchSize", self._options.pop("batchSize", None))
+
+        # If the cursor option is already specified, avoid overriding it.
+        self._options.setdefault("cursor", {})
+        # If the pipeline performs a write, we ignore the initial batchSize
+        # since the server doesn't return results in this case.
+        if self._batch_size is not None and not self._performs_write:
+            self._options["cursor"]["batchSize"] = self._batch_size
+
         self._cursor_class = cursor_class
         self._explicit_session = explicit_session
         self._user_fields = user_fields
@@ -51,11 +66,6 @@ def __init__(self, target, cursor_class, pipeline, options,
             options.pop('collation', None))
 
         self._max_await_time_ms = options.pop('maxAwaitTimeMS', None)
-        self._batch_size = common.validate_non_negative_integer_or_none(
-            "batchSize", options.pop("batchSize", None))
-
-        self._dollar_out = (self._pipeline and
-                            '$out' in self._pipeline[-1])
 
     @property
     def _aggregation_target(self):
@@ -99,16 +109,17 @@ def get_cursor(self, session, server, sock_info, slave_ok):
         # - server version is >= 4.2 or
         # - server version is >= 3.2 and pipeline doesn't use $out
         if (('readConcern' not in cmd) and
-                ((sock_info.max_wire_version >= 4 and not self._dollar_out) or
+                ((sock_info.max_wire_version >= 4 and
+                  not self._performs_write) or
                  (sock_info.max_wire_version >= 8))):
             read_concern = self._target.read_concern
         else:
             read_concern = None
 
         # Apply this target's write concern if:
         # writeConcern has not been specified as a kwarg and pipeline doesn't
-        # use $out
-        if 'writeConcern' not in cmd and self._dollar_out:
+        # perform a write operation
+        if 'writeConcern' not in cmd and self._performs_write:
             write_concern = self._target._write_concern_for(session)
         else:
             write_concern = None
@@ -159,6 +170,16 @@ def get_cursor(self, session, server, sock_info, slave_ok):
 
 
 class _CollectionAggregationCommand(_AggregationCommand):
+    def __init__(self, *args, **kwargs):
+        # Pop additional option and initialize parent class.
+        use_cursor = kwargs.pop("use_cursor", True)
+        super(_CollectionAggregationCommand, self).__init__(*args, **kwargs)
+
+        # Remove the cursor document if the user has set use_cursor to False.
+        self._use_cursor = use_cursor
+        if not self._use_cursor:
+            self._options.pop("cursor", None)
+
     @property
     def _aggregation_target(self):
         return self._target.name
@@ -172,6 +193,15 @@ def _database(self):
         return self._target.database
 
 
+class _CollectionRawAggregationCommand(_CollectionAggregationCommand):
+    def __init__(self, *args, **kwargs):
+        super(_CollectionRawAggregationCommand, self).__init__(*args, **kwargs)
+
+        # For raw-batches, we set the initial batchSize for the cursor to 0.
+        if self._use_cursor and not self._performs_write:
+            self._options["cursor"]["batchSize"] = 0
+
+
 class _DatabaseAggregationCommand(_AggregationCommand):
     @property
     def _aggregation_target(self):
diff --git a/pymongo/change_stream.py b/pymongo/change_stream.py
@@ -98,6 +98,7 @@ def _client(self):
         raise NotImplementedError
 
     def _change_stream_options(self):
+        """Return the options dict for the $changeStream pipeline stage."""
         options = {}
         if self._full_document is not None:
             options['fullDocument'] = self._full_document
@@ -110,6 +111,7 @@ def _change_stream_options(self):
         return options
 
     def _command_options(self):
+        """Return the options dict for the aggregation command."""
         options = {'cursor': {}}
         if self._max_await_time_ms is not None:
             options["maxAwaitTimeMS"] = self._max_await_time_ms
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -29,7 +29,8 @@
 from pymongo import (common,
                      helpers,
                      message)
-from pymongo.aggregation import _CollectionAggregationCommand
+from pymongo.aggregation import (_CollectionAggregationCommand,
+                                 _CollectionRawAggregationCommand)
 from pymongo.bulk import BulkOperationBuilder, _Bulk
 from pymongo.command_cursor import CommandCursor, RawBatchCommandCursor
 from pymongo.common import ORDERED_TYPES
@@ -2258,11 +2259,8 @@ def options(self, session=None):
 
         return options
 
-    def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
+    def _aggregate(self, aggregation_command, pipeline, cursor_class, session,
                    explicit_session, **kwargs):
-        # Check if we use the $out stage
-        dollar_out = pipeline and '$out' in pipeline[-1]
-
         # Remove things that are not command options.
         use_cursor = True
         if "useCursor" in kwargs:
@@ -2271,25 +2269,14 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
                 "and will be removed in PyMongo 4.0",
                 DeprecationWarning, stacklevel=2)
             use_cursor = common.validate_boolean(
-                "useCursor", kwargs.pop("useCursor"))
-
-        # If the server does not support the "cursor" option we
-        # ignore useCursor and batchSize.
-        if use_cursor:
-            if "cursor" not in kwargs:
-                kwargs["cursor"] = {}
-            # Ignore batchSize when the $out pipeline stage is used.
-            # batchSize is meaningless in that case since the server
-            # doesn't return results. This also avoids SERVER-23923.
-            if first_batch_size is not None and not dollar_out:
-                kwargs["cursor"]["batchSize"] = first_batch_size
-
-        cmd = _CollectionAggregationCommand(
+                "useCursor", kwargs.pop("useCursor", True))
+
+        cmd = aggregation_command(
             self, cursor_class, pipeline, kwargs, explicit_session,
-            user_fields={'cursor': {'firstBatch': 1}})
+            user_fields={'cursor': {'firstBatch': 1}}, use_cursor=use_cursor)
         return self.__database.client._retryable_read(
             cmd.get_cursor, self._read_preference_for(session), session,
-            retryable=not dollar_out)
+            retryable=not cmd._performs_write)
 
     def aggregate(self, pipeline, session=None, **kwargs):
         """Perform an aggregation using the aggregation framework on this
@@ -2314,11 +2301,11 @@ def aggregate(self, pipeline, session=None, **kwargs):
           - `useCursor` (bool): Deprecated. Will be removed in PyMongo 4.0.
 
         The :meth:`aggregate` method obeys the :attr:`read_preference` of this
-        :class:`Collection`. Please note that using the ``$out`` pipeline stage
-        requires a read preference of
+        :class:`Collection`. Please note that using the ``$out`` and ``$merge``
+        pipeline stages requires a read preference of
         :attr:`~pymongo.read_preferences.ReadPreference.PRIMARY` (the default).
-        The server will raise an error if the ``$out`` pipeline stage is used
-        with any other read preference.
+        The server will raise an error if the ``$out`` or ``$merge`` pipeline
+        stages are used with any other read preference.
 
         .. note:: This method does not support the 'explain' option. Please
            use :meth:`~pymongo.database.Database.command` instead. An
@@ -2338,6 +2325,8 @@ def aggregate(self, pipeline, session=None, **kwargs):
           A :class:`~pymongo.command_cursor.CommandCursor` over the result
           set.
 
+        .. versionchanged:: 3.9
+           Added support for the ``$merge`` pipeline stage.
         .. versionchanged:: 3.9
            Apply this collection's read concern to pipelines containing the
            `$out` stage when connected to MongoDB >= 4.2.
@@ -2364,9 +2353,9 @@ def aggregate(self, pipeline, session=None, **kwargs):
             https://docs.mongodb.com/manual/reference/command/aggregate
         """
         with self.__database.client._tmp_session(session, close=False) as s:
-            return self._aggregate(pipeline,
+            return self._aggregate(_CollectionAggregationCommand,
+                                   pipeline,
                                    CommandCursor,
-                                   kwargs.get('batchSize'),
                                    session=s,
                                    explicit_session=session is not None,
                                    **kwargs)
@@ -2397,8 +2386,13 @@ def aggregate_raw_batches(self, pipeline, **kwargs):
         if "session" in kwargs:
             raise ConfigurationError(
                 "aggregate_raw_batches does not support sessions")
-        return self._aggregate(pipeline, RawBatchCommandCursor, 0,
-                               None, False, **kwargs)
+
+        return self._aggregate(_CollectionRawAggregationCommand,
+                               pipeline,
+                               RawBatchCommandCursor,
+                               session=None,
+                               explicit_session=False,
+                               **kwargs)
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -446,29 +446,6 @@ def _fix_outgoing(self, son, collection):
             son = manipulator.transform_outgoing(son, collection)
         return son
 
-    def _aggregate(self, pipeline, cursor_class, session,
-                   explicit_session, **kwargs):
-        # Check if we use the $out stage
-        dollar_out = pipeline and '$out' in pipeline[-1]
-
-        # Specify cursor option if it has not been provided by the user.
-        if "cursor" not in kwargs:
-            kwargs["cursor"] = {}
-
-        # Ignore batchSize when the $out pipeline stage is used.
-        # batchSize is meaningless in that case since the server
-        # doesn't return results. This also avoids SERVER-23923.
-        batch_size = kwargs.get("batchSize")
-        if batch_size is not None and not dollar_out:
-            kwargs["cursor"]["batchSize"] = batch_size
-
-        cmd = _DatabaseAggregationCommand(
-            self, cursor_class, pipeline, kwargs, explicit_session,
-            user_fields={'cursor': {'firstBatch': 1}})
-        return self.client._retryable_read(
-            cmd.get_cursor, self._read_preference_for(session), session,
-            retryable=not dollar_out)
-
     def aggregate(self, pipeline, session=None, **kwargs):
         """Perform a database-level aggregation.
 
@@ -500,11 +477,11 @@ def aggregate(self, pipeline, session=None, **kwargs):
             :class:`~pymongo.collation.Collation`.
 
         The :meth:`aggregate` method obeys the :attr:`read_preference` of this
-        :class:`Database`. Please note that using the ``$out`` pipeline stage
-        requires a read preference of
+        :class:`Database`. Please note that using the ``$out`` or ``$merge``
+        pipeline stages requires a read preference of
         :attr:`~pymongo.read_preferences.ReadPreference.PRIMARY` (the default).
-        The server will raise an error if the ``$out`` pipeline stage is used
-        with any other read preference.
+        The server will raise an error if the ``$out`` or ``$merge`` pipeline
+        stages is used with any other read preference.
 
         .. note:: This method does not support the 'explain' option. Please
            use :meth:`~pymongo.database.Database.command` instead.
@@ -531,11 +508,12 @@ def aggregate(self, pipeline, session=None, **kwargs):
             https://docs.mongodb.com/manual/reference/command/aggregate
         """
         with self.client._tmp_session(session, close=False) as s:
-            return self._aggregate(pipeline,
-                                   CommandCursor,
-                                   session=s,
-                                   explicit_session=session is not None,
-                                   **kwargs)
+            cmd = _DatabaseAggregationCommand(
+                self, CommandCursor, pipeline, kwargs, session is not None,
+                user_fields={'cursor': {'firstBatch': 1}})
+            return self.client._retryable_read(
+                cmd.get_cursor, self._read_preference_for(s), s,
+                retryable=not cmd._performs_write)
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
diff --git a/test/crud/v2/aggregate-merge.json b/test/crud/v2/aggregate-merge.json
