Add support for maxAwaitTimeMS to getMore

Author: aherlihy

pymongo/cursor.py

@@ -154,6 +154,7 @@ def __init__(self, collection, filter=None, projection=None, skip=0,
  self.__hint = None
  self.__comment = None
  self.__max_time_ms = None
+ self.__max_await_time_ms = None
  self.__max = None
  self.__min = None
  self.__manipulate = manipulate
@@ -244,10 +245,10 @@ def _clone(self, deepcopy=True):
  """Internal clone helper."""
  clone = self._clone_base()
  values_to_clone = ("spec", "projection", "skip", "limit",
- "max_time_ms", "comment", "max", "min",
- "ordering", "explain", "hint", "batch_size",
- "max_scan", "manipulate", "query_flags",
- "modifiers")
+ "max_time_ms", "max_await_time_ms", "comment",
+ "max", "min", "ordering", "explain", "hint",
+ "batch_size", "max_scan", "manipulate",
+ "query_flags", "modifiers")
  data = dict((k, v) for k, v in iteritems(self.__dict__)
  if k.startswith('_Cursor__') and k[9:] in values_to_clone)
  if deepcopy:
@@ -470,6 +471,35 @@ def max_time_ms(self, max_time_ms):
  self.__max_time_ms = max_time_ms
  return self
 
+ def max_await_time_ms(self, max_await_time_ms):
+ """Specifies a time limit for a getMore operation on a
+ :attr:`~pymongo.CursorType.TAILABLE_AWAIT` cursor. For all other types
+ of cursor max_await_time_ms is ignored.
+
+ Raises :exc:`TypeError` if `max_await_time_ms` is not an integer or
+ ``None``.
+ Raises :exc:`~pymongo.errors.InvalidOperation` if this :class:`Cursor`
+ has already been used.
+
+ .. note:: `max_await_time_ms` requires server version **>= 3.2**
+
+ :Parameters:
+ - `max_await_time_ms`: the time limit after which the operation is
+ aborted
+
+ .. versionadded:: 3.2
+ """
+ if (not isinstance(max_await_time_ms, integer_types)
+ and max_await_time_ms is not None):
+ raise TypeError("max_await_time_ms must be an integer or None")
+ self.__check_okay_to_chain()
+
+ # Ignore max_await_time_ms if not tailable or await_data is False.
+ if self.__query_flags & CursorType.TAILABLE_AWAIT:
+ self.__max_await_time_ms = max_await_time_ms
+
+ return self
+
  def __getitem__(self, index):
  """Get a single document or a slice of documents from this cursor.
 
@@ -1007,7 +1037,7 @@ def _refresh(self):
  limit,
  self.__id,
  self.__codec_options,
- self.__max_time_ms))
+ self.__max_await_time_ms))
 
  else: # Cursor id is zero nothing else to return
  self.__killed = True

pymongo/message.py

@@ -196,14 +196,14 @@ def _gen_find_command(coll, spec, projection, skip, limit, batch_size,
  return cmd
 
 
-def _gen_get_more_command(cursor_id, coll, batch_size, max_time_ms):
+def _gen_get_more_command(cursor_id, coll, batch_size, max_await_time_ms):
  """Generate a getMore command document."""
  cmd = SON([('getMore', cursor_id),
  ('collection', coll)])
  if batch_size:
  cmd['batchSize'] = batch_size
- if max_time_ms:
- cmd['maxTimeMS'] = max_time_ms
+ if max_await_time_ms is not None:
+ cmd['maxTimeMS'] = max_await_time_ms
  return cmd
 
 
@@ -274,24 +274,25 @@ def get_message(self, set_slave_ok, is_mongos, use_cmd=False):
 class _GetMore(object):
  """A getmore operation."""
 
- __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_time_ms',
+ __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_await_time_ms',
  'codec_options')
 
  name = 'getMore'
 
  def __init__(self, db, coll, ntoreturn, cursor_id, codec_options,
- max_time_ms=None):
+ max_await_time_ms=None):
  self.db = db
  self.coll = coll
  self.ntoreturn = ntoreturn
  self.cursor_id = cursor_id
  self.codec_options = codec_options
- self.max_time_ms = max_time_ms
+ self.max_await_time_ms = max_await_time_ms
 
  def as_command(self):
  """Return a getMore command document for this query."""
  return _gen_get_more_command(self.cursor_id, self.coll,
- self.ntoreturn, self.max_time_ms), self.db
+ self.ntoreturn,
+ self.max_await_time_ms), self.db
 
  def get_message(self, dummy0, dummy1, use_cmd=False):
  """Get a getmore message."""

test/test_cursor.py

@@ -25,6 +25,7 @@
 from bson.py3compat import u, PY3
 from bson.son import SON
 from pymongo import (MongoClient,
+ monitoring,
  ASCENDING,
  DESCENDING,
  ALL,
@@ -41,7 +42,7 @@
  host,
  port,
  IntegrationTest)
-from test.utils import server_started_with_auth
+from test.utils import server_started_with_auth, single_client, EventListener
 
 if PY3:
  long = int
@@ -190,6 +191,114 @@ def test_max_time_ms(self):
  "maxTimeAlwaysTimeOut",
  mode="off")
 
+ @client_context.require_version_min(3, 1, 9, -1)
+ def test_max_await_time_ms(self):
+ db = self.db
+ db.pymongo_test.drop()
+ coll = db.create_collection("pymongo_test", capped=True, size=4096)
+
+ self.assertRaises(TypeError, coll.find().max_await_time_ms, 'foo')
+ coll.insert_one({"amalia": 1})
+ coll.insert_one({"amalia": 2})
+
+ coll.find().max_await_time_ms(None)
+ coll.find().max_await_time_ms(long(1))
+
+ # When cursor is not tailable_await
+ cursor = coll.find()
+ self.assertEqual(None, cursor._Cursor__max_await_time_ms)
+ cursor = coll.find().max_await_time_ms(99)
+ self.assertEqual(None, cursor._Cursor__max_await_time_ms)
+
+ # If cursor is tailable_await and timeout is unset
+ cursor = coll.find(cursor_type=CursorType.TAILABLE_AWAIT)
+ self.assertEqual(None, cursor._Cursor__max_await_time_ms)
+
+ # If cursor is tailable_await and timeout is set
+ cursor = coll.find(
+ cursor_type=CursorType.TAILABLE_AWAIT).max_await_time_ms(99)
+ self.assertEqual(99, cursor._Cursor__max_await_time_ms)
+
+ cursor = coll.find(
+ cursor_type=CursorType.TAILABLE_AWAIT).max_await_time_ms(
+ 10).max_await_time_ms(90)
+ self.assertEqual(90, cursor._Cursor__max_await_time_ms)
+
+ listener = EventListener()
+ saved_listeners = monitoring._LISTENERS
+ monitoring._LISTENERS = monitoring._Listeners([])
+ coll = single_client(
+ event_listeners=[listener])[self.db.name].pymongo_test
+ results = listener.results
+
+ try:
+ # Tailable_await defaults.
+ list(coll.find(cursor_type=CursorType.TAILABLE_AWAIT))
+ # find
+ self.assertFalse('maxTimeMS' in results['started'][0].command)
+ # getMore
+ self.assertFalse('maxTimeMS' in results['started'][1].command)
+ results.clear()
+
+ # Tailable_await with max_await_time_ms set.
+ list(coll.find(
+ cursor_type=CursorType.TAILABLE_AWAIT).max_await_time_ms(99))
+ # find
+ self.assertFalse('maxTimeMS' in results['started'][0].command)
+ # getMore
+ self.assertTrue('maxTimeMS' in results['started'][1].command)
+ self.assertEqual(99, results['started'][1].command['maxTimeMS'])
+ results.clear()
+
+ # Tailable_await with max_time_ms
+ list(coll.find(
+ cursor_type=CursorType.TAILABLE_AWAIT).max_time_ms(1))
+ # find
+ self.assertTrue('maxTimeMS' in results['started'][0].command)
+ self.assertEqual(1, results['started'][0].command['maxTimeMS'])
+ # getMore
+ self.assertFalse('maxTimeMS' in results['started'][1].command)
+ results.clear()
+
+ # Tailable_await with both max_time_ms and max_await_time_ms
+ list(coll.find(
+ cursor_type=CursorType.TAILABLE_AWAIT).max_time_ms(
+ 1).max_await_time_ms(99))
+ # find
+ self.assertTrue('maxTimeMS' in results['started'][0].command)
+ self.assertEqual(1, results['started'][0].command['maxTimeMS'])
+ # getMore
+ self.assertTrue('maxTimeMS' in results['started'][1].command)
+ self.assertEqual(99, results['started'][1].command['maxTimeMS'])
+ results.clear()
+
+ # Non tailable_await with max_await_time_ms
+ list(coll.find(batch_size=1).max_await_time_ms(99))
+ # find
+ self.assertFalse('maxTimeMS' in results['started'][0].command)
+ # getMore
+ self.assertFalse('maxTimeMS' in results['started'][1].command)
+ results.clear()
+
+ # Non tailable_await with max_time_ms
+ list(coll.find(batch_size=1).max_time_ms(99))
+ # find
+ self.assertTrue('maxTimeMS' in results['started'][0].command)
+ self.assertEqual(99, results['started'][0].command['maxTimeMS'])
+ # getMore
+ self.assertFalse('maxTimeMS' in results['started'][1].command)
+
+ # Non tailable_await with both max_time_ms and max_await_time_ms
+ list(coll.find(batch_size=1).max_time_ms(99).max_await_time_ms(88))
+ # find
+ self.assertTrue('maxTimeMS' in results['started'][0].command)
+ self.assertEqual(99, results['started'][0].command['maxTimeMS'])
+ # getMore
+ self.assertFalse('maxTimeMS' in results['started'][1].command)
+
+ finally:
+ monitoring._LISTENERS = saved_listeners
+
  @client_context.require_version_min(2, 5, 3, -1)
  @client_context.require_test_commands
  def test_max_time_ms_getmore(self):