PYTHON-1720 Add start_after parameter to watch() methods

ShaneHarvey · ShaneHarvey · commit afbf18b0ad30 · 2019-06-04T16:06:47.000-07:00
diff --git a/pymongo/change_stream.py b/pymongo/change_stream.py
@@ -52,7 +52,7 @@ class ChangeStream(object):
     """
     def __init__(self, target, pipeline, full_document, resume_after,
                  max_await_time_ms, batch_size, collation,
-                 start_at_operation_time, session):
+                 start_at_operation_time, session, start_after):
         if pipeline is None:
             pipeline = []
         elif not isinstance(pipeline, list):
@@ -82,6 +82,7 @@ def __init__(self, target, pipeline, full_document, resume_after,
         self._collation = collation
         self._start_at_operation_time = start_at_operation_time
         self._session = session
+        self._start_after = copy.deepcopy(start_after)
         self._cursor = self._create_cursor()
 
     @property
@@ -101,6 +102,8 @@ def _pipeline_options(self):
             options['fullDocument'] = self._full_document
         if self._resume_token is not None:
             options['resumeAfter'] = self._resume_token
+        if self._start_after is not None:
+            options['startAfter'] = self._start_after
         if self._start_at_operation_time is not None:
             options['startAtOperationTime'] = self._start_at_operation_time
         return options
@@ -141,7 +144,7 @@ def _cmd(session, server, sock_info, slave_ok):
 
             if (self._start_at_operation_time is None and
                 self._resume_token is None and
-                cursor.get("_id") is None and
+                self._start_after is None and
                 sock_info.max_wire_version >= 7):
                 self._start_at_operation_time = result["operationTime"]
 
@@ -277,6 +280,7 @@ def try_next(self):
                 "token is missing.")
         self._resume_token = copy.copy(resume_token)
         self._start_at_operation_time = None
+        self._start_after = None
 
         if self._decode_custom:
             return _bson_to_dict(change.raw, self._orig_codec_options)
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -2458,7 +2458,7 @@ def aggregate_raw_batches(self, pipeline, **kwargs):
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
-              start_at_operation_time=None, session=None):
+              start_at_operation_time=None, session=None, start_after=None):
         """Watch changes on this collection.
 
         Performs an aggregation with an implicit initial ``$changeStream``
@@ -2519,8 +2519,10 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             updates will include both a delta describing the changes to the
             document, as well as a copy of the entire document that was
             changed from some time after the change occurred.
-          - `resume_after` (optional): The logical starting point for this
-            change stream.
+          - `resume_after` (optional): A resume token. If provided, the
+            change stream will start returning changes that occur directly
+            after the operation specified in the resume token. A resume token
+            is the _id value of a change document.
           - `max_await_time_ms` (optional): The maximum time in milliseconds
             for the server to wait for changes before responding to a getMore
             operation.
@@ -2534,10 +2536,16 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             MongoDB >= 4.0.
           - `session` (optional): a
             :class:`~pymongo.client_session.ClientSession`.
+          - `start_after` (optional): The same as `resume_after` except that
+            `start_after` can resume notifications after an invalidate event.
+            This option and `resume_after` are mutually exclusive.
 
         :Returns:
           A :class:`~pymongo.change_stream.CollectionChangeStream` cursor.
 
+        .. versionchanged:: 3.9
+           Added the ``start_after`` parameter.
+
         .. versionchanged:: 3.7
            Added the ``start_at_operation_time`` parameter.
 
@@ -2550,8 +2558,8 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
         """
         return CollectionChangeStream(
             self, pipeline, full_document, resume_after, max_await_time_ms,
-            batch_size, collation, start_at_operation_time, session
-        )
+            batch_size, collation, start_at_operation_time, session,
+            start_after)
 
     def group(self, key, condition, initial, reduce, finalize=None, **kwargs):
         """Perform a query similar to an SQL *group by* operation.
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -447,7 +447,7 @@ def _fix_outgoing(self, son, collection):
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
-              start_at_operation_time=None, session=None):
+              start_at_operation_time=None, session=None, start_after=None):
         """Watch changes on this database.
 
         Performs an aggregation with an implicit initial ``$changeStream``
@@ -499,8 +499,10 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             updates will include both a delta describing the changes to the
             document, as well as a copy of the entire document that was
             changed from some time after the change occurred.
-          - `resume_after` (optional): The logical starting point for this
-            change stream.
+          - `resume_after` (optional): A resume token. If provided, the
+            change stream will start returning changes that occur directly
+            after the operation specified in the resume token. A resume token
+            is the _id value of a change document.
           - `max_await_time_ms` (optional): The maximum time in milliseconds
             for the server to wait for changes before responding to a getMore
             operation.
@@ -514,10 +516,16 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             MongoDB >= 4.0.
           - `session` (optional): a
             :class:`~pymongo.client_session.ClientSession`.
+          - `start_after` (optional): The same as `resume_after` except that
+            `start_after` can resume notifications after an invalidate event.
+            This option and `resume_after` are mutually exclusive.
 
         :Returns:
           A :class:`~pymongo.change_stream.DatabaseChangeStream` cursor.
 
+        .. versionchanged:: 3.9
+           Added the ``start_after`` parameter.
+
         .. versionadded:: 3.7
 
         .. mongodoc:: changeStreams
@@ -527,8 +535,8 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
         """
         return DatabaseChangeStream(
             self, pipeline, full_document, resume_after, max_await_time_ms,
-            batch_size, collation, start_at_operation_time, session
-        )
+            batch_size, collation, start_at_operation_time, session,
+            start_after)
 
     def _command(self, sock_info, command, slave_ok=False, value=1, check=True,
                  allowable_errors=None, read_preference=ReadPreference.PRIMARY,
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -802,7 +802,7 @@ def _server_property(self, attr_name):
 
     def watch(self, pipeline=None, full_document='default', resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
-              start_at_operation_time=None, session=None):
+              start_at_operation_time=None, session=None, start_after=None):
         """Watch changes on this cluster.
 
         Performs an aggregation with an implicit initial ``$changeStream``
@@ -854,8 +854,10 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             updates will include both a delta describing the changes to the
             document, as well as a copy of the entire document that was
             changed from some time after the change occurred.
-          - `resume_after` (optional): The logical starting point for this
-            change stream.
+          - `resume_after` (optional): A resume token. If provided, the
+            change stream will start returning changes that occur directly
+            after the operation specified in the resume token. A resume token
+            is the _id value of a change document.
           - `max_await_time_ms` (optional): The maximum time in milliseconds
             for the server to wait for changes before responding to a getMore
             operation.
@@ -869,10 +871,16 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
             MongoDB >= 4.0.
           - `session` (optional): a
             :class:`~pymongo.client_session.ClientSession`.
+          - `start_after` (optional): The same as `resume_after` except that
+            `start_after` can resume notifications after an invalidate event.
+            This option and `resume_after` are mutually exclusive.
 
         :Returns:
           A :class:`~pymongo.change_stream.ClusterChangeStream` cursor.
 
+        .. versionchanged:: 3.9
+           Added the ``start_after`` parameter.
+
         .. versionadded:: 3.7
 
         .. mongodoc:: changeStreams
@@ -882,8 +890,8 @@ def watch(self, pipeline=None, full_document='default', resume_after=None,
         """
         return ClusterChangeStream(
             self.admin, pipeline, full_document, resume_after, max_await_time_ms,
-            batch_size, collation, start_at_operation_time, session
-        )
+            batch_size, collation, start_at_operation_time, session,
+            start_after)
 
     @property
     def event_listeners(self):
diff --git a/test/test_change_stream.py b/test/test_change_stream.py
@@ -597,6 +597,64 @@ def test_read_concern(self):
         with coll.watch():
             pass
 
+    def invalidate_resume_token(self):
+        with self.coll.watch(
+                [{'$match': {'operationType': 'invalidate'}}]) as cs:
+            self.coll.insert_one({'_id': 1})
+            self.coll.drop()
+            resume_token = cs.next()['_id']
+            self.assertFalse(cs.alive)
+            return resume_token
+
+    @client_context.require_version_min(4, 1, 1)
+    def test_start_after(self):
+        resume_token = self.invalidate_resume_token()
+
+        # resume_after cannot resume after invalidate.
+        with self.assertRaises(OperationFailure):
+            self.coll.watch(resume_after=resume_token)
+
+        # start_after can resume after invalidate.
+        with self.coll.watch(start_after=resume_token) as change_stream:
+            self.coll.insert_one({'_id': 2})
+            change = change_stream.next()
+            self.assertEqual(change['operationType'], 'insert')
+            self.assertEqual(change['fullDocument'], {'_id': 2})
+
+    @client_context.require_version_min(4, 1, 1)
+    def test_start_after_resume_process_with_changes(self):
+        resume_token = self.invalidate_resume_token()
+
+        with self.coll.watch(start_after=resume_token,
+                             max_await_time_ms=250) as change_stream:
+            self.coll.insert_one({'_id': 2})
+            change = change_stream.next()
+            self.assertEqual(change['operationType'], 'insert')
+            self.assertEqual(change['fullDocument'], {'_id': 2})
+
+            self.assertIsNone(change_stream.try_next())
+            self.kill_change_stream_cursor(change_stream)
+
+            self.coll.insert_one({'_id': 3})
+            change = change_stream.next()
+            self.assertEqual(change['operationType'], 'insert')
+            self.assertEqual(change['fullDocument'], {'_id': 3})
+
+    @client_context.require_no_mongos  # Remove after SERVER-41196
+    @client_context.require_version_min(4, 1, 1)
+    def test_start_after_resume_process_without_changes(self):
+        resume_token = self.invalidate_resume_token()
+
+        with self.coll.watch(start_after=resume_token,
+                             max_await_time_ms=250) as change_stream:
+            self.assertIsNone(change_stream.try_next())
+            self.kill_change_stream_cursor(change_stream)
+
+            self.coll.insert_one({'_id': 2})
+            change = change_stream.next()
+            self.assertEqual(change['operationType'], 'insert')
+            self.assertEqual(change['fullDocument'], {'_id': 2})
+
 
 class TestAllScenarios(unittest.TestCase):
 
