PYTHON-894 - Set cursor.alive False after final batch.

Even with this change, "next" can raise StopIteration even though
"alive" is True. For example if batch size is 2 and there are 4
documents in the result set, then after the 4th document "alive" is True
but "next" raises StopIteration.

Author: ajdavis

pymongo/command_cursor.py

@@ -40,7 +40,7 @@ def __init__(self, collection, cursor_info,
         )
         self.__retrieved = retrieved
         self.__batch_size = 0
-        self.__killed = False
+        self.__killed = (self.__id == 0)
 
         if "ns" in cursor_info:
             self.__ns = cursor_info["ns"]
@@ -123,6 +123,8 @@ def __send_message(self, msg):
             client.disconnect()
             raise
         self.__id = response["cursor_id"]
+        if self.__id == 0:
+            self.__killed = True
 
         assert response["starting_from"] == self.__retrieved, (
             "Result batch started from %s, expected %s" % (
@@ -153,7 +155,14 @@ def _refresh(self):
 
     @property
     def alive(self):
-        """Does this cursor have the potential to return more data?"""
+        """Does this cursor have the potential to return more data?
+
+        Even if :attr:`alive` is ``True``, :meth:`.next` can raise
+        :exc:`StopIteration`. Best to use a for loop::
+
+            for doc in collection.aggregate(pipeline):
+                print(doc)
+        """
         return bool(len(self.__data) or (not self.__killed))
 
     @property

pymongo/cursor.py

@@ -977,6 +977,8 @@ def __send_message(self, message):
             raise
 
         self.__id = response["cursor_id"]
+        if self.__id == 0:
+            self.__killed = True
 
         # starting from doesn't get set on getmore's for tailable cursors
         if not (self.__query_flags & _QUERY_OPTIONS["tailable_cursor"]):
@@ -1050,6 +1052,11 @@ def alive(self):
         since they will stop iterating even though they *may* return more
         results in the future.
 
+        With regular cursors, simply use a for loop instead of :attr:`alive`::
+
+            for doc in collection.find():
+                print(doc)
+
         .. versionadded:: 1.5
         """
         return bool(len(self.__data) or (not self.__killed))

test/test_collection.py

@@ -1430,6 +1430,23 @@ def test_aggregation_cursor(self):
                 expected_sum,
                 sum(doc['_id'] for doc in cursor))
 
+    def test_aggregation_cursor_alive(self):
+        if not version.at_least(self.db.connection, (2, 5, 1)):
+            raise SkipTest("Aggregation cursor requires MongoDB >= 2.5.1")
+        self.db.test.remove()
+        self.db.test.insert([{} for _ in range(3)])
+        self.addCleanup(self.db.test.remove)
+        cursor = self.db.test.aggregate(pipeline=[], cursor={'batchSize': 2})
+        n = 0
+        while True:
+            cursor.next()
+            n += 1
+            if 3 == n:
+                self.assertFalse(cursor.alive)
+                break
+
+            self.assertTrue(cursor.alive)
+
     def test_parallel_scan(self):
         if is_mongos(self.db.connection):
             raise SkipTest("mongos does not support parallel_scan")

test/test_cursor.py

@@ -1165,5 +1165,20 @@ def close(self, dummy):
             client.set_cursor_manager(CursorManager)
             ctx.exit()
 
+    def test_alive(self):
+        self.db.test.remove()
+        self.db.test.insert([{} for _ in range(3)])
+        self.addCleanup(self.db.test.remove)
+        cursor = self.db.test.find().batch_size(2)
+        n = 0
+        while True:
+            cursor.next()
+            n += 1
+            if 3 == n:
+                self.assertFalse(cursor.alive)
+                break
+
+            self.assertTrue(cursor.alive)
+
 if __name__ == "__main__":
     unittest.main()