Merge branch 'cursor-maxTimeMS'

* cursor-maxTimeMS: Add support for maxTimeMS option in cursors
doctrine · Jan 6, 2016 · 7b4d519 · 7b4d519
2 parents 4d3f855 + e459513
commit 7b4d519
Show file tree

Hide file tree

Showing 8 changed files with 119 additions and 1 deletion.
diff --git a/lib/Doctrine/MongoDB/Cursor.php b/lib/Doctrine/MongoDB/Cursor.php
@@ -68,6 +68,7 @@ class Cursor implements CursorInterface
     protected $options = array();
     protected $batchSize;
     protected $limit;
+    protected $maxTimeMS;
     protected $readPreference;
     protected $readPreferenceTags;
     protected $skip;
@@ -445,6 +446,20 @@ public function limit($num)
         return $this;
     }
 
+    /**
+     * Wrapper method for MongoCursor::maxTimeMS().
+     *
+     * @see http://php.net/manual/en/mongocursor.maxtimems.php
+     * @param integer $ms
+     * @return self
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->maxTimeMS = (integer) $ms;
+        $this->mongoCursor->maxTimeMS($this->maxTimeMS);
+        return $this;
+    }
+
     /**
      * Wrapper method for MongoCursor::next().
      *
@@ -480,6 +495,9 @@ public function recreate()
         if ($this->limit !== null) {
             $this->mongoCursor->limit($this->limit);
         }
+        if ($this->maxTimeMS !== null) {
+            $this->mongoCursor->maxTimeMS($this->maxTimeMS);
+        }
         if ($this->skip !== null) {
             $this->mongoCursor->skip($this->skip);
         }

diff --git a/lib/Doctrine/MongoDB/EagerCursor.php b/lib/Doctrine/MongoDB/EagerCursor.php
@@ -362,6 +362,20 @@ public function limit($num)
         return $this;
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function maxTimeMS($ms)
+    {
+        // Need to use method_exists - adding to CursorInterface is not allowed
+        // due to SemVer restrictions
+        if (method_exists($this->cursor, 'maxTimeMS')) {
+            $this->cursor->maxTimeMS($ms);
+        }
+
+        return $this;
+    }
+
     /**
      * {@inheritDoc}
      */

diff --git a/lib/Doctrine/MongoDB/LoggableCursor.php b/lib/Doctrine/MongoDB/LoggableCursor.php
@@ -102,6 +102,19 @@ public function limit($num)
         return parent::limit($num);
     }
 
+    /**
+     * @see Cursor::maxTimeMS()
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->log(array(
+            'maxTimeMS' => true,
+            'maxTimeMSNum' => $ms,
+        ));
+
+        return parent::maxTimeMS($ms);
+    }
+
     /**
      * @see Cursor::skip()
      */

diff --git a/lib/Doctrine/MongoDB/Query/Builder.php b/lib/Doctrine/MongoDB/Query/Builder.php
@@ -1072,6 +1072,18 @@ public function maxDistance($maxDistance)
         return $this;
     }
 
+    /**
+     * Specifies a cumulative time limit in milliseconds for processing operations on a cursor.
+     *
+     * @param int $ms
+     * @return $this
+     */
+    public function maxTimeMS($ms)
+    {
+        $this->query['maxTimeMS'] = $ms;
+        return $this;
+    }
+
     /**
      * Updates the value of the field to a specified value if the specified value is less than the current value of the field.
      *

diff --git a/lib/Doctrine/MongoDB/Query/Query.php b/lib/Doctrine/MongoDB/Query/Query.php
@@ -395,7 +395,7 @@ protected function prepareCursor(Cursor $cursor)
             $cursor->setReadPreference($this->query['readPreference'], $this->query['readPreferenceTags']);
         }
 
-        foreach ($this->getQueryOptions('hint', 'immortal', 'limit', 'skip', 'slaveOkay', 'sort') as $key => $value) {
+        foreach ($this->getQueryOptions('hint', 'immortal', 'limit', 'maxTimeMS', 'skip', 'slaveOkay', 'sort') as $key => $value) {
             $cursor->$key($value);
         }
 

diff --git a/tests/Doctrine/MongoDB/Tests/CursorTest.php b/tests/Doctrine/MongoDB/Tests/CursorTest.php
@@ -319,6 +319,40 @@ public function testRecreate()
         $cursor->recreate();
     }
 
+    public function testSetMaxTimeMSWhenRecreateCursor()
+    {
+        $self = $this;
+
+        $setCursorExpectations = function($mongoCursor) use ($self) {
+            $mongoCursor->expects($self->once())
+                ->method('maxTimeMS')
+                ->with(30000);
+        };
+
+        $mongoCursor = $this->getMockMongoCursor();
+        $recreatedMongoCursor = $this->getMockMongoCursor();
+
+        $setCursorExpectations($mongoCursor);
+        $setCursorExpectations($recreatedMongoCursor);
+
+        $mongoCollection = $this->getMockCollection();
+        $mongoCollection->expects($this->once())
+            ->method('find')
+            ->with(array('x' => 9500), array())
+            ->will($this->returnValue($recreatedMongoCursor));
+
+        $collection = $this->getMockCollection();
+        $collection->expects($this->once())
+            ->method('getMongoCollection')
+            ->will($this->returnValue($mongoCollection));
+
+        $cursor = $this->getTestCursor($collection, $mongoCursor, array('x' => 9500));
+
+        $cursor->maxTimeMS(30000);
+
+        $cursor->recreate();
+    }
+
     private function getMockCollection()
     {
         return $this->getMockBuilder('Doctrine\MongoDB\Collection')

diff --git a/tests/Doctrine/MongoDB/Tests/LoggableCursorTest.php b/tests/Doctrine/MongoDB/Tests/LoggableCursorTest.php
@@ -23,6 +23,7 @@ public function provideLoggedMethods()
             array('limit'),
             array('hint', array()),
             array('snapshot'),
+            array('maxTimeMS', array())
         );
     }
 }
diff --git a/tests/Doctrine/MongoDB/Tests/Query/QueryTest.php b/tests/Doctrine/MongoDB/Tests/Query/QueryTest.php
@@ -286,6 +286,32 @@ public function testUseIdentifierKeys()
         $this->assertSame($cursor, $query->execute());
     }
 
+    public function testSpecifyMaxTimeMSOnCursor()
+    {
+        $cursor = $this->getMockCursor();
+        $collection = $this->getMockCollection();
+
+        $collection->expects($this->once())
+            ->method('find')
+            ->with($this->equalTo(array('foo' => 'bar')))
+            ->will($this->returnValue($cursor));
+
+        $cursor->expects($this->once())
+            ->method('maxTimeMS')
+            ->with($this->equalTo(30000))
+            ->will($this->returnValue($cursor));
+
+        $queryArray = array(
+            'type' => Query::TYPE_FIND,
+            'query' => array('foo' => 'bar'),
+            'maxTimeMS' => 30000
+        );
+
+        $query = new Query($collection, $queryArray, array());
+
+        $this->assertSame($cursor, $query->execute());
+    }
+
     /**
      * @return \Doctrine\MongoDB\Collection
      */