[Issue 634] Merged changes that fixes issue MorphiaOrg#634. Added uni…

…t test for the changes, and backed out the changed to countAll that are not necessary.

morphia/src/main/java/org/mongodb/morphia/query/Query.java

@@ -7,6 +7,7 @@
 import org.bson.types.CodeWScope;
 
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 
 /**
@@ -93,6 +94,14 @@ public interface Query<T> extends QueryResults<T>, Cloneable {
    */
   Query<T> maxScan(int value);
 
+  /**
+   * Specifies a time limit for executing the query. Requires server version 2.6 or above.
+   *
+   * @param maxTime must be > 0.  A value < 0 indicates no limit
+   * @param maxTimeUnit
+   */
+  Query<T> maxTime(long maxTime, TimeUnit maxTimeUnit);
+
   /**
    * This makes it possible to attach a comment to a query. Because these comments propagate to the profile log, adding comments
    * can make your profile data much easier to interpret and trace.
@@ -289,4 +298,4 @@ public interface Query<T> extends QueryResults<T>, Cloneable {
    * Creates and returns a copy of this {@link Query}.
    */
   Query<T> cloneQuery();
-}
+}

morphia/src/main/java/org/mongodb/morphia/query/QueryImpl.java

@@ -26,6 +26,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import static java.lang.String.format;
 import static org.mongodb.morphia.query.QueryValidator.validateQuery;
@@ -63,6 +64,8 @@ public class QueryImpl<T> extends CriteriaContainerImpl implements Query<T> {
     private boolean tailAwaitData;
     private ReadPreference readPref;
     private Integer maxScan;
+    private Long maxTime;
+    private TimeUnit maxTimeUnit;
     private String comment;
     private boolean returnKey;
 
@@ -252,6 +255,10 @@ public DBCursor prepareCursor() {
             cursor.addSpecial("$maxScan", maxScan);
         }
 
+        if (maxTime != null && maxTimeUnit != null) {
+            cursor.maxTime(maxTime, maxTimeUnit);
+        }
+
         if (max != null) {
             cursor.addSpecial("$max", max);
         }
@@ -420,6 +427,12 @@ public Query<T> maxScan(final int value) {
         return this;
     }
 
+    public Query<T> maxTime(final long value, final TimeUnit timeUnitValue) {
+        maxTime = value > 0 ? value : null;
+        maxTimeUnit = timeUnitValue;
+        return this;
+    }
+
     public Query<T> comment(final String comment) {
         this.comment = comment;
         return this;
@@ -617,4 +630,4 @@ public Map<String, Object> explain() {
         DBCursor cursor = prepareCursor();
         return (BasicDBObject) cursor.explain();
     }
-}
+}

morphia/src/test/java/com/mongodb/StubDBCollection.java

@@ -0,0 +1,85 @@
+package com.mongodb;
+
+import java.util.List;
+
+/**
+ * Somewhat ugly, but effective, way to stub out the collection.  Still needs a real DB, otherwise you get null pointers, 
+ * but it does allow you to override with your own behaviours.  
+ * 
+ * This is currently being used to unit test QueryImpl.
+ */
+public class StubDBCollection extends DBCollection {
+
+    private StubDBCursor stubDBCursor;
+
+    public StubDBCollection(final DB db) {
+        super(db, "");
+    }
+
+    @Override
+    public WriteResult insert(final List<DBObject> list, final WriteConcern concern, final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    public WriteResult update(final DBObject q, final DBObject o, final boolean upsert, final boolean multi,
+                              final WriteConcern concern, final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    protected void doapply(final DBObject o) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    public WriteResult remove(final DBObject o, final WriteConcern concern, final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    QueryResultIterator find(final DBObject ref, final DBObject fields, final int numToSkip, final int batchSize, final int limit,
+                             final int options, final ReadPreference readPref,
+                             final DBDecoder decoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    QueryResultIterator find(final DBObject ref, final DBObject fields, final int numToSkip, final int batchSize, final int limit,
+                             final int options, final ReadPreference readPref,
+                             final DBDecoder decoder, final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    @Deprecated
+    public void createIndex(final DBObject keys, final DBObject options, final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    public Cursor aggregate(final List<DBObject> pipeline, final AggregationOptions options, final ReadPreference readPreference) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    public List<Cursor> parallelScan(final ParallelScanOptions options) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    BulkWriteResult executeBulkWriteOperation(final boolean ordered, final List<WriteRequest> requests, final WriteConcern writeConcern,
+                                              final DBEncoder encoder) {
+        throw new UnsupportedOperationException("Not implemented yet!");
+    }
+
+    @Override
+    public DBCursor find(final DBObject ref, final DBObject keys) {
+        return stubDBCursor; 
+    }
+
+    public void setDBCursor(final StubDBCursor stubDBCursor) {
+        this.stubDBCursor = stubDBCursor;
+    }
+
+}

morphia/src/test/java/com/mongodb/StubDBCursor.java

@@ -0,0 +1,33 @@
+package com.mongodb;
+
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Somewhat ugly, but effective, way to stub out the collection.  Can be supplied with a StubDBCollection. Can override the cursor 
+ * methods to check what's actually being called.
+ *
+ * This is currently being used to unit test QueryImpl.
+ */
+public class StubDBCursor extends DBCursor {
+    private long maxTime;
+    private TimeUnit maxTimeUnit;
+
+    public StubDBCursor(final DBCollection collection) {
+        super(collection, new BasicDBObject(), new BasicDBObject(), ReadPreference.primary());
+    }
+
+    @Override
+    public DBCursor maxTime(final long maxTime, final TimeUnit timeUnit) {
+        this.maxTime = maxTime;
+        this.maxTimeUnit = timeUnit;
+        return this;
+    }
+
+    public long getMaxTime() {
+        return maxTime;
+    }
+
+    public TimeUnit getMaxTimeUnit() {
+        return maxTimeUnit;
+    }
+}

morphia/src/test/java/org/mongodb/morphia/query/QueryImplTest.java

@@ -0,0 +1,37 @@
+package org.mongodb.morphia.query;
+
+import com.mongodb.StubDBCollection;
+import com.mongodb.StubDBCursor;
+import org.junit.Test;
+import org.mongodb.morphia.TestBase;
+import org.mongodb.morphia.entities.SimpleEntity;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.junit.Assert.assertThat;
+
+public class QueryImplTest extends TestBase {
+
+    @Test
+    public void testMaxTimeApplied() {
+        // given
+        StubDBCollection coll = new StubDBCollection(getDb());
+        StubDBCursor stubDBCursor = new StubDBCursor(coll);
+        coll.setDBCursor(stubDBCursor);
+
+        QueryImpl<SimpleEntity> query = new QueryImpl<SimpleEntity>(SimpleEntity.class, coll, getDs());
+
+        long maxTime = 123;
+        TimeUnit maxTimeUnit = TimeUnit.MILLISECONDS;
+
+        // when
+        query.maxTime(maxTime, maxTimeUnit);
+        query.prepareCursor();
+
+        // then
+        assertThat(stubDBCursor.getMaxTime(), is(maxTime));
+        assertThat(stubDBCursor.getMaxTimeUnit(), is(maxTimeUnit));
+    }
+
+}