DATADOC-141 DATADOC-89 introduced a protected prepareCollection method and a boolean flag setSlaveOk on the MongoTemplate to have template instance level control over slave behavior

spring-data-mongodb/src/main/java/org/springframework/data/document/mongodb/MongoTemplate.java

@@ -93,6 +93,11 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 	 */
 	private WriteResultChecking writeResultChecking = WriteResultChecking.NONE;
 
+	/*
+	 * Flag used to indicate use of slaveOk() for any operations on collections.
+	 */
+	private boolean slaveOk = false;
+
 	private final MongoConverter mongoConverter;
 	private final MappingContext<? extends MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 	private final MongoDbFactory mongoDbFactory;
@@ -248,85 +253,13 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 		Assert.notNull(callback);
 
 		try {
-			DBCollection collection = getDb().getCollection(collectionName);
+			DBCollection collection = getAndPrepareCollection(getDb(), collectionName);
 			return callback.doInCollection(collection);
 		} catch (RuntimeException e) {
 			throw potentiallyConvertRuntimeException(e);
 		}
 	}
 
-	/**
-	 * Central callback executing method to do queries against the datastore that requires reading a single object from a
-	 * collection of objects. It will take the following steps
-	 * <ol>
-	 * <li>Execute the given {@link ConnectionCallback} for a {@link DBObject}.</li>
-	 * <li>Apply the given {@link DbObjectCallback} to each of the {@link DBObject}s to obtain the result.</li>
-	 * <ol>
-	 * 
-	 * @param <T>
-	 * @param collectionCallback
-	 *          the callback to retrieve the {@link DBObject} with
-	 * @param objectCallback
-	 *          the {@link DbObjectCallback} to transform {@link DBObject}s into the actual domain type
-	 * @param collectionName
-	 *          the collection to be queried
-	 * @return
-	 */
-	private <T> T execute(CollectionCallback<DBObject> collectionCallback, DbObjectCallback<T> objectCallback,
-			String collectionName) {
-
-		try {
-			T result = objectCallback.doWith(collectionCallback.doInCollection(getCollection(collectionName)));
-			return result;
-		} catch (RuntimeException e) {
-			throw potentiallyConvertRuntimeException(e);
-		}
-	}
-
-	/**
-	 * Central callback executing method to do queries against the datastore that requires reading a collection of
-	 * objects. It will take the following steps
-	 * <ol>
-	 * <li>Execute the given {@link ConnectionCallback} for a {@link DBCursor}.</li>
-	 * <li>Prepare that {@link DBCursor} with the given {@link CursorPreparer} (will be skipped if {@link CursorPreparer}
-	 * is {@literal null}</li>
-	 * <li>Iterate over the {@link DBCursor} and applies the given {@link DbObjectCallback} to each of the
-	 * {@link DBObject}s collecting the actual result {@link List}.</li>
-	 * <ol>
-	 * 
-	 * @param <T>
-	 * @param collectionCallback
-	 *          the callback to retrieve the {@link DBCursor} with
-	 * @param preparer
-	 *          the {@link CursorPreparer} to potentially modify the {@link DBCursor} before ireating over it
-	 * @param objectCallback
-	 *          the {@link DbObjectCallback} to transform {@link DBObject}s into the actual domain type
-	 * @param collectionName
-	 *          the collection to be queried
-	 * @return
-	 */
-	private <T> List<T> executeEach(CollectionCallback<DBCursor> collectionCallback, CursorPreparer preparer,
-			DbObjectCallback<T> objectCallback, String collectionName) {
-
-		try {
-			DBCursor cursor = collectionCallback.doInCollection(getCollection(collectionName));
-
-			if (preparer != null) {
-				cursor = preparer.prepare(cursor);
-			}
-
-			List<T> result = new ArrayList<T>();
-
-			for (DBObject object : cursor) {
-				result.add(objectCallback.doWith(object));
-			}
-
-			return result;
-		} catch (RuntimeException e) {
-			throw potentiallyConvertRuntimeException(e);
-		}
-	}
-
 	/* (non-Javadoc)
 	 * @see org.springframework.data.document.mongodb.MongoOperations#executeInSession(org.springframework.data.document.mongodb.DBCallback)
 	 */
@@ -485,6 +418,18 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 	public void insert(String collectionName, Object objectToSave) {
 		doInsert(collectionName, objectToSave, this.mongoConverter);
 	}
+	
+	/**
+	 * Prepare the collection before any processing is done using it. This allows a convenient way to apply 
+	 * settings like slaveOk() etc. Can be overridden in sub-classes.
+	 * 
+	 * @param collection
+	 */
+	protected void prepareCollection(DBCollection collection) {
+		if(this.slaveOk) {
+			collection.slaveOk();
+		}
+	}
 
 	protected <T> void doInsert(String collectionName, T objectToSave, MongoWriter<T> writer) {
 		BasicDBObject dbDoc = new BasicDBObject();
@@ -823,12 +768,12 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 	 * @see org.springframework.data.document.mongodb.MongoOperations#getCollection(java.lang.Class)
 	 */
 	public <T> List<T> getCollection(Class<T> targetClass) {
-		return executeEach(new FindCallback(null), null, new ReadDbObjectCallback<T>(mongoConverter, targetClass),
+		return executeFindMultiInternal(new FindCallback(null), null, new ReadDbObjectCallback<T>(mongoConverter, targetClass),
 				determineCollectionName(targetClass));
 	}
 
 	public <T> List<T> getCollection(String collectionName, Class<T> targetClass) {
-		return executeEach(new FindCallback(null), null, new ReadDbObjectCallback<T>(mongoConverter, targetClass),
+		return executeFindMultiInternal(new FindCallback(null), null, new ReadDbObjectCallback<T>(mongoConverter, targetClass),
 				collectionName);
 	}
 
@@ -887,7 +832,7 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(targetClass);
 		DBObject mappedQuery = mapper.getMappedObject(query, entity);
 
-		return execute(new FindOneCallback(mappedQuery, fields), new ReadDbObjectCallback<T>(readerToUse, targetClass),
+		return executeFindOneInternal(new FindOneCallback(mappedQuery, fields), new ReadDbObjectCallback<T>(readerToUse, targetClass),
 				collectionName);
 	}
 
@@ -921,7 +866,7 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 			LOGGER.debug("find using query: " + query + " fields: " + fields + " for class: " + targetClass
 					+ " in collection: " + collectionName);
 		}
-		return executeEach(new FindCallback(mapper.getMappedObject(query, entity), fields), preparer,
+		return executeFindMultiInternal(new FindCallback(mapper.getMappedObject(query, entity), fields), preparer,
 				new ReadDbObjectCallback<T>(mongoConverter, targetClass), collectionName);
 	}
 
@@ -949,7 +894,7 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 		}
 		MongoReader<? super T> readerToUse = this.mongoConverter;
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(targetClass);
-		return executeEach(new FindCallback(mapper.getMappedObject(query, entity), fields), null,
+		return executeFindMultiInternal(new FindCallback(mapper.getMappedObject(query, entity), fields), null,
 				new ReadDbObjectCallback<T>(readerToUse, targetClass), collectionName);
 	}
 
@@ -993,7 +938,7 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 					+ targetClass + " in collection: " + collectionName);
 		}
 		MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(targetClass);
-		return execute(new FindAndRemoveCallback(mapper.getMappedObject(query, entity), fields, sort),
+		return executeFindOneInternal(new FindAndRemoveCallback(mapper.getMappedObject(query, entity), fields, sort),
 				new ReadDbObjectCallback<T>(readerToUse, targetClass), collectionName);
 	}
 
@@ -1043,6 +988,88 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 		}
 	}
 
+	private DBCollection getAndPrepareCollection(DB db, String collectionName) {
+		try {
+			DBCollection collection = db.getCollection(collectionName);
+			prepareCollection(collection);
+			return collection;
+		} catch (RuntimeException e) {
+			throw potentiallyConvertRuntimeException(e);
+		}
+	}
+
+	/**
+	 * Internal method using callbacks to do queries against the datastore that requires reading a single object from a
+	 * collection of objects. It will take the following steps
+	 * <ol>
+	 * <li>Execute the given {@link ConnectionCallback} for a {@link DBObject}.</li>
+	 * <li>Apply the given {@link DbObjectCallback} to each of the {@link DBObject}s to obtain the result.</li>
+	 * <ol>
+	 * 
+	 * @param <T>
+	 * @param collectionCallback
+	 *          the callback to retrieve the {@link DBObject} with
+	 * @param objectCallback
+	 *          the {@link DbObjectCallback} to transform {@link DBObject}s into the actual domain type
+	 * @param collectionName
+	 *          the collection to be queried
+	 * @return
+	 */
+	private <T> T executeFindOneInternal(CollectionCallback<DBObject> collectionCallback, DbObjectCallback<T> objectCallback,
+			String collectionName) {
+
+		try {
+			T result = objectCallback.doWith(collectionCallback.doInCollection(getAndPrepareCollection(getDb(), collectionName)));
+			return result;
+		} catch (RuntimeException e) {
+			throw potentiallyConvertRuntimeException(e);
+		}
+	}
+
+	/**
+	 * Internal method using callback to do queries against the datastore that requires reading a collection of
+	 * objects. It will take the following steps
+	 * <ol>
+	 * <li>Execute the given {@link ConnectionCallback} for a {@link DBCursor}.</li>
+	 * <li>Prepare that {@link DBCursor} with the given {@link CursorPreparer} (will be skipped if {@link CursorPreparer}
+	 * is {@literal null}</li>
+	 * <li>Iterate over the {@link DBCursor} and applies the given {@link DbObjectCallback} to each of the
+	 * {@link DBObject}s collecting the actual result {@link List}.</li>
+	 * <ol>
+	 * 
+	 * @param <T>
+	 * @param collectionCallback
+	 *          the callback to retrieve the {@link DBCursor} with
+	 * @param preparer
+	 *          the {@link CursorPreparer} to potentially modify the {@link DBCursor} before ireating over it
+	 * @param objectCallback
+	 *          the {@link DbObjectCallback} to transform {@link DBObject}s into the actual domain type
+	 * @param collectionName
+	 *          the collection to be queried
+	 * @return
+	 */
+	private <T> List<T> executeFindMultiInternal(CollectionCallback<DBCursor> collectionCallback, CursorPreparer preparer,
+			DbObjectCallback<T> objectCallback, String collectionName) {
+
+		try {
+			DBCursor cursor = collectionCallback.doInCollection(getAndPrepareCollection(getDb(), collectionName));
+
+			if (preparer != null) {
+				cursor = preparer.prepare(cursor);
+			}
+
+			List<T> result = new ArrayList<T>();
+
+			for (DBObject object : cursor) {
+				result.add(objectCallback.doWith(object));
+			}
+
+			return result;
+		} catch (RuntimeException e) {
+			throw potentiallyConvertRuntimeException(e);
+		}
+	}
+
 	private MongoPersistentEntity<?> getPersistentEntity(Class<?> type) {
 		return type == null ? null : mappingContext.getPersistentEntity(type);
 	}
@@ -1272,4 +1299,8 @@ public class MongoTemplate implements MongoOperations, ApplicationEventPublisher
 		this.writeConcern = writeConcern;
 	}
 
+	public void setSlaveOk(boolean slaveOk) {
+		this.slaveOk = slaveOk;
+	}
+
 }

spring-data-mongodb/src/test/java/org/springframework/data/document/mongodb/MongoTemplateTests.java

@@ -37,6 +37,7 @@ import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.dao.DataAccessException;
 import org.springframework.dao.DataIntegrityViolationException;
 import org.springframework.data.document.InvalidDocumentStoreApiUsageException;
 import org.springframework.data.document.mongodb.convert.MappingMongoConverter;
@@ -54,6 +55,7 @@ import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import com.mongodb.DBCollection;
 import com.mongodb.DBObject;
 import com.mongodb.Mongo;
+import com.mongodb.MongoException;
 import com.mongodb.WriteResult;
 
 /**
@@ -512,4 +514,26 @@ public class MongoTemplateTests {
 
 	}
 
+	@Test
+	public void testUsingSlaveOk() throws Exception {
+		this.template.execute("slaveOkTest", new CollectionCallback<Object>() {
+			public Object doInCollection(DBCollection collection)
+					throws MongoException, DataAccessException {
+				assertThat(collection.getOptions(), is(0));
+				assertThat(collection.getDB().getOptions(), is(0));
+				return null;
+			}
+		});
+		MongoTemplate slaveTemplate = new MongoTemplate(this.template.getDbFactory());
+		slaveTemplate.setSlaveOk(true);
+		slaveTemplate.execute("slaveOkTest", new CollectionCallback<Object>() {
+			public Object doInCollection(DBCollection collection)
+					throws MongoException, DataAccessException {
+				assertThat(collection.getOptions(), is(4));
+				assertThat(collection.getDB().getOptions(), is(0));
+				return null;
+			}
+		});
+	}
+
 }