PYTHON-2389 Add session support to find_raw_batches and aggregate_raw_batches (#658)

(cherry picked from commit 0e0c4fd)

Author: prashantmital

pymongo/collection.py

@@ -1547,21 +1547,16 @@ def find_raw_batches(self, *args, **kwargs):
           >>> for batch in cursor:
           ...     print(bson.decode_all(batch))
 
-        .. note:: find_raw_batches does not support sessions or auto
-           encryption.
+        .. note:: find_raw_batches does not support auto encryption.
 
         .. versionchanged:: 3.12
            Instead of ignoring the user-specified read concern, this method
            now sends it to the server when connected to MongoDB 3.6+.
 
+           Added session support.
+
         .. versionadded:: 3.6
         """
-        # OP_MSG with document stream returns is required to support
-        # sessions.
-        if "session" in kwargs:
-            raise ConfigurationError(
-                "find_raw_batches does not support sessions")
-
         # OP_MSG is required to support encryption.
         if self.__database.client._encrypter:
             raise InvalidOperation(
@@ -2505,7 +2500,7 @@ def aggregate(self, pipeline, session=None, **kwargs):
                                    explicit_session=session is not None,
                                    **kwargs)
 
-    def aggregate_raw_batches(self, pipeline, **kwargs):
+    def aggregate_raw_batches(self, pipeline, session=None, **kwargs):
         """Perform an aggregation and retrieve batches of raw BSON.
 
         Similar to the :meth:`aggregate` method but returns a
@@ -2522,28 +2517,25 @@ def aggregate_raw_batches(self, pipeline, **kwargs):
           >>> for batch in cursor:
           ...     print(bson.decode_all(batch))
 
-        .. note:: aggregate_raw_batches does not support sessions or auto
-           encryption.
+        .. note:: aggregate_raw_batches does not support auto encryption.
+
+        .. versionchanged:: 3.12
+           Added session support.
 
         .. versionadded:: 3.6
         """
-        # OP_MSG with document stream returns is required to support
-        # sessions.
-        if "session" in kwargs:
-            raise ConfigurationError(
-                "aggregate_raw_batches does not support sessions")
-
         # OP_MSG is required to support encryption.
         if self.__database.client._encrypter:
             raise InvalidOperation(
                 "aggregate_raw_batches does not support auto encryption")
 
-        return self._aggregate(_CollectionRawAggregationCommand,
-                               pipeline,
-                               RawBatchCommandCursor,
-                               session=None,
-                               explicit_session=False,
-                               **kwargs)
+        with self.__database.client._tmp_session(session, close=False) as s:
+            return self._aggregate(_CollectionRawAggregationCommand,
+                                   pipeline,
+                                   RawBatchCommandCursor,
+                                   session=s,
+                                   explicit_session=session is not None,
+                                   **kwargs)
 
     def watch(self, pipeline=None, full_document=None, resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,

pymongo/message.py

@@ -473,6 +473,8 @@ def use_command(self, sock_info):
 
 class _RawBatchGetMore(_GetMore):
     def use_command(self, sock_info):
+        # Compatibility checks.
+        super(_RawBatchGetMore, self).use_command(sock_info)
         if sock_info.max_wire_version >= 8:
             # MongoDB 4.2+ supports exhaust over OP_MSG
             return True

test/test_cursor.py

@@ -46,6 +46,7 @@
                   unittest,
                   IntegrationTest)
 from test.utils import (EventListener,
+                        OvertCommandListener,
                         ignore_deprecations,
                         rs_or_single_client,
                         WhiteListEventListener)
@@ -1478,6 +1479,76 @@ def test_manipulate(self):
         with self.assertRaises(InvalidOperation):
             c.find_raw_batches(manipulate=True)
 
+    @client_context.require_transactions
+    def test_find_raw_transaction(self):
+        c = self.db.test
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener])
+        with client.start_session() as session:
+            with session.start_transaction():
+                batches = list(client[self.db.name].test.find_raw_batches(
+                    session=session).sort('_id'))
+                cmd = listener.results['started'][0]
+                self.assertEqual(cmd.command_name, 'find')
+                self.assertEqual(cmd.command['$clusterTime'],
+                                 decode_all(session.cluster_time.raw)[0])
+                self.assertEqual(cmd.command['startTransaction'], True)
+                self.assertEqual(cmd.command['txnNumber'], 1)
+
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+
+    @client_context.require_sessions
+    @client_context.require_failCommand_fail_point
+    def test_find_raw_retryable_reads(self):
+        c = self.db.test
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener],
+                                     retryReads=True)
+        with self.fail_point({
+            'mode': {'times': 1}, 'data': {'failCommands': ['find'],
+                                           'closeConnection': True}}):
+            batches = list(
+                client[self.db.name].test.find_raw_batches().sort('_id'))
+
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+        self.assertEqual(len(listener.results['started']), 2)
+        for cmd in listener.results['started']:
+            self.assertEqual(cmd.command_name, 'find')
+
+    @client_context.require_version_min(5, 0, 0)
+    @client_context.require_no_standalone
+    def test_find_raw_snapshot_reads(self):
+        c = self.db.get_collection(
+            "test", write_concern=WriteConcern(w="majority"))
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener],
+                                     retryReads=True)
+        db = client[self.db.name]
+        with client.start_session(snapshot=True) as session:
+            db.test.distinct('x', {}, session=session)
+            batches = list(db.test.find_raw_batches(
+                session=session).sort('_id'))
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+
+        find_cmd = listener.results['started'][1].command
+        self.assertEqual(find_cmd['readConcern']['level'], 'snapshot')
+        self.assertIsNotNone(find_cmd['readConcern']['atClusterTime'])
+
     def test_explain(self):
         c = self.db.test
         c.insert_one({})
@@ -1602,6 +1673,75 @@ def test_aggregate_raw(self):
         self.assertEqual(1, len(batches))
         self.assertEqual(docs, decode_all(batches[0]))
 
+    @client_context.require_transactions
+    def test_aggregate_raw_transaction(self):
+        c = self.db.test
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener])
+        with client.start_session() as session:
+            with session.start_transaction():
+                batches = list(client[self.db.name].test.aggregate_raw_batches(
+                    [{'$sort': {'_id': 1}}], session=session))
+                cmd = listener.results['started'][0]
+                self.assertEqual(cmd.command_name, 'aggregate')
+                self.assertEqual(cmd.command['$clusterTime'], session.cluster_time)
+                self.assertEqual(cmd.command['startTransaction'], True)
+                self.assertEqual(cmd.command['txnNumber'], 1)
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+
+    @client_context.require_sessions
+    @client_context.require_failCommand_fail_point
+    def test_aggregate_raw_retryable_reads(self):
+        c = self.db.test
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener],
+                                     retryReads=True)
+        with self.fail_point({
+            'mode': {'times': 1}, 'data': {'failCommands': ['aggregate'],
+                                           'closeConnection': True}}):
+            batches = list(client[self.db.name].test.aggregate_raw_batches(
+                [{'$sort': {'_id': 1}}]))
+
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+        self.assertEqual(len(listener.results['started']), 3)
+        cmds = listener.results['started']
+        self.assertEqual(cmds[0].command_name, 'aggregate')
+        self.assertEqual(cmds[1].command_name, 'aggregate')
+
+    @client_context.require_version_min(5, 0, -1)
+    @client_context.require_no_standalone
+    def test_aggregate_raw_snapshot_reads(self):
+        c = self.db.get_collection(
+            "test", write_concern=WriteConcern(w="majority"))
+        c.drop()
+        docs = [{'_id': i, 'x': 3.0 * i} for i in range(10)]
+        c.insert_many(docs)
+
+        listener = OvertCommandListener()
+        client = rs_or_single_client(event_listeners=[listener],
+                                     retryReads=True)
+        db = client[self.db.name]
+        with client.start_session(snapshot=True) as session:
+            db.test.distinct('x', {}, session=session)
+            batches = list(db.test.aggregate_raw_batches(
+                [{'$sort': {'_id': 1}}], session=session))
+        self.assertEqual(1, len(batches))
+        self.assertEqual(docs, decode_all(batches[0]))
+
+        find_cmd = listener.results['started'][1].command
+        self.assertEqual(find_cmd['readConcern']['level'], 'snapshot')
+        self.assertIsNotNone(find_cmd['readConcern']['atClusterTime'])
+
     def test_server_error(self):
         c = self.db.test
         c.drop()

test/test_session.py

@@ -891,6 +891,12 @@ def test_reads(self):
             lambda coll, session: coll.count_documents({}, session=session))
         self._test_reads(
             lambda coll, session: coll.distinct('foo', session=session))
+        self._test_reads(
+            lambda coll, session: list(coll.aggregate_raw_batches(
+                [], session=session)))
+        self._test_reads(
+            lambda coll, session: list(coll.find_raw_batches(
+                {}, session=session)))
 
         # SERVER-40938 removed support for casually consistent mapReduce.
         map_reduce_exc = None
@@ -916,16 +922,6 @@ def scan(coll, session):
             self._test_reads(
                 lambda coll, session: scan(coll, session=session))
 
-        self.assertRaises(
-            ConfigurationError,
-            self._test_reads,
-            lambda coll, session: list(
-                coll.aggregate_raw_batches([], session=session)))
-        self.assertRaises(
-            ConfigurationError,
-            self._test_reads,
-            lambda coll, session: list(
-                coll.find_raw_batches({}, session=session)))
         self.assertRaises(
             ConfigurationError,
             self._test_reads,