Updates to MongoTemplate to use a Mongo instance to lazily create DB instances rather than directly using a DB instance. Also integrated with Spring TransactionSynchronization so a single DB instance can be shared by multiple templates in a single thread

spring-datastore-mongodb/src/main/java/org/springframework/datastore/document/mongodb/DBHolder.java

@@ -0,0 +1,71 @@
+package org.springframework.datastore.document.mongodb;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.springframework.transaction.support.ResourceHolderSupport;
+import org.springframework.util.Assert;
+
+import com.mongodb.DB;
+
+public class DBHolder extends ResourceHolderSupport {
+	private static final Object DEFAULT_KEY = new Object();
+
+	private final Map<Object, DB> dbMap = new ConcurrentHashMap<Object, DB>();
+
+
+	public DBHolder(DB db) {
+		addDB(db);
+	}
+
+	public DBHolder(Object key, DB db) {
+		addDB(key, db);
+	}
+
+
+	public DB getDB() {
+		return getDB(DEFAULT_KEY);
+	}
+
+	public DB getDB(Object key) {
+		return this.dbMap.get(key);
+	}
+
+
+	public DB getAnyDB() {
+		if (!this.dbMap.isEmpty()) {
+			return this.dbMap.values().iterator().next();
+		}
+		return null;
+	}
+
+	public void addDB(DB session) {
+		addDB(DEFAULT_KEY, session);
+	}
+
+	public void addDB(Object key, DB session) {
+		Assert.notNull(key, "Key must not be null");
+		Assert.notNull(session, "DB must not be null");
+		this.dbMap.put(key, session);
+	}
+
+	public DB removeDB(Object key) {
+		return this.dbMap.remove(key);
+	}
+
+	public boolean containsDB(DB session) {
+		return this.dbMap.containsValue(session);
+	}
+
+	public boolean isEmpty() {
+		return this.dbMap.isEmpty();
+	}
+
+	public boolean doesNotHoldNonDefaultDB() {
+		synchronized (this.dbMap) {
+			return this.dbMap.isEmpty() ||
+					(this.dbMap.size() == 1 && this.dbMap.containsKey(DEFAULT_KEY));
+		}
+	}
+
+}

spring-datastore-mongodb/src/main/java/org/springframework/datastore/document/mongodb/MongoDbUtils.java

@@ -16,10 +16,20 @@
 
 package org.springframework.datastore.document.mongodb;
 
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.springframework.dao.DataAccessException;
+import org.springframework.dao.DataAccessResourceFailureException;
+import org.springframework.dao.DataIntegrityViolationException;
 import org.springframework.datastore.document.UncategorizedDocumentStoreException;
+import org.springframework.transaction.support.TransactionSynchronizationManager;
+import org.springframework.util.Assert;
 
+import com.mongodb.DB;
+import com.mongodb.Mongo;
 import com.mongodb.MongoException;
+import com.mongodb.MongoException.DuplicateKey;
+import com.mongodb.MongoException.Network;
 
 /**
  * Helper class featuring helper methods for internal MongoDb classes.
@@ -27,10 +37,14 @@ import com.mongodb.MongoException;
  * <p>Mainly intended for internal use within the framework.
  *
  * @author Thomas Risberg
+ * @author Graeme Rocher
+ * 
  * @since 1.0
  */
 public class MongoDbUtils {
 
+	static final Log logger = LogFactory.getLog(MongoDbUtils.class);
+	
 	/**
 	 * Convert the given runtime exception to an appropriate exception from the
 	 * <code>org.springframework.dao</code> hierarchy.
@@ -45,6 +59,12 @@ public class MongoDbUtils {
 		// Check for well-known MongoException subclasses.
 		
 		// All other MongoExceptions
+		if(ex instanceof DuplicateKey) {
+			return new DataIntegrityViolationException(ex.getMessage(),ex);
+		}
+		if(ex instanceof Network) {
+			return new DataAccessResourceFailureException(ex.getMessage(), ex);
+		}
 		if (ex instanceof MongoException) {
 			return new UncategorizedDocumentStoreException(ex.getMessage(), ex);
 		}
@@ -55,4 +75,94 @@ public class MongoDbUtils {
 		return null;				
 	}
 
+	public static DB getDB(Mongo mongo, String databaseName) {
+		return doGetDB(mongo, databaseName, true);
+	}
+
+	public static DB doGetDB(Mongo mongo, String databaseName, boolean allowCreate) {
+		Assert.notNull(mongo, "No Mongo instance specified");
+
+		DBHolder dbHolder = (DBHolder) TransactionSynchronizationManager.getResource(mongo);
+		if (dbHolder != null && !dbHolder.isEmpty()) {
+			// pre-bound Mongo DB 
+			DB db = null;
+			if (TransactionSynchronizationManager.isSynchronizationActive() &&
+					dbHolder.doesNotHoldNonDefaultDB()) {
+				// Spring transaction management is active ->
+				db = dbHolder.getDB();
+				if (db != null && !dbHolder.isSynchronizedWithTransaction()) {
+					logger.debug("Registering Spring transaction synchronization for existing Mongo DB");
+					TransactionSynchronizationManager.registerSynchronization(new MongoSynchronization(dbHolder, mongo));
+					dbHolder.setSynchronizedWithTransaction(true);
+				}
+			}
+			if (db != null) {
+				return db;
+			}
+		}
+
+		logger.debug("Opening Mongo DB");
+		DB db = mongo.getDB(databaseName);
+		// Use same Session for further Mongo actions within the transaction.
+		// Thread object will get removed by synchronization at transaction completion.
+		if (TransactionSynchronizationManager.isSynchronizationActive()) {
+			// We're within a Spring-managed transaction, possibly from JtaTransactionManager.
+			logger.debug("Registering Spring transaction synchronization for new Hibernate Session");
+			DBHolder holderToUse = dbHolder;
+			if (holderToUse == null) {
+				holderToUse = new DBHolder(db);
+			}
+			else {
+				holderToUse.addDB(db);
+			}
+			TransactionSynchronizationManager.registerSynchronization(new MongoSynchronization(dbHolder, mongo));
+			holderToUse.setSynchronizedWithTransaction(true);
+			if (holderToUse != dbHolder) {
+				TransactionSynchronizationManager.bindResource(mongo, holderToUse);
+			}
+		}
+
+		// Check whether we are allowed to return the DB.
+		if (!allowCreate && !isDBTransactional(db, mongo)) {
+			throw new IllegalStateException("No Mongo DB bound to thread, " +
+			    "and configuration does not allow creation of non-transactional one here");
+		}
+
+		return db;
+	}
+	
+
+	/**
+	 * Return whether the given DB instance is transactional, that is,
+	 * bound to the current thread by Spring's transaction facilities.
+	 * @param db the DB to check
+	 * @param mongo the Mongo instance that the DB was created with
+	 * (may be <code>null</code>)
+	 * @return whether the DB is transactional
+	 */
+	public static boolean isDBTransactional(DB db, Mongo mongo) {
+		if (mongo == null) {
+			return false;
+		}
+		DBHolder dbHolder =
+				(DBHolder) TransactionSynchronizationManager.getResource(mongo);
+		return (dbHolder != null && dbHolder.containsDB(db));
+	}
+	
+	/**
+	 * Perform actual closing of the Mongo DB object,
+	 * catching and logging any cleanup exceptions thrown.
+	 * @param db the DB to close (may be <code>null</code>)
+	 */
+	public static void closeDB(DB db) {
+		if (db != null) {
+			logger.debug("Closing Mongo DB object");
+			try {
+				db.requestDone();
+			}
+			catch (Throwable ex) {
+				logger.debug("Unexpected exception on closing Mongo DB object", ex);
+			}
+		}
+	}	
 }

spring-datastore-mongodb/src/main/java/org/springframework/datastore/document/mongodb/MongoFactoryBean.java

@@ -24,6 +24,7 @@ import org.springframework.util.Assert;
 
 import com.mongodb.DB;
 import com.mongodb.Mongo;
+import com.mongodb.MongoOptions;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -32,9 +33,11 @@ import org.apache.commons.logging.LogFactory;
  * Convenient factory for configuring MongoDB.
  *
  * @author Thomas Risberg
+ * @author Graeme Rocher
+ * 
  * @since 1.0
  */
-public class MongoDbFactoryBean implements FactoryBean<DB>, InitializingBean, 
+public class MongoFactoryBean implements FactoryBean<Mongo>, InitializingBean, 
 		PersistenceExceptionTranslator {
 
 	/**
@@ -43,14 +46,20 @@ public class MongoDbFactoryBean implements FactoryBean<DB>, InitializingBean,
 	protected final Log logger = LogFactory.getLog(getClass());
 
 	private Mongo mongo;
+	private MongoOptions mongoOptions;
 	private String host;
 	private Integer port;
 	private String databaseName;
 	
+	
 	public void setMongo(Mongo mongo) {
 		this.mongo = mongo;
 	}
 
+	public void setMongoOptions(MongoOptions mongoOptions) {
+		this.mongoOptions = mongoOptions;
+	}
+
 	public void setDatabaseName(String databaseName) {
 		this.databaseName = databaseName;
 	}
@@ -63,14 +72,13 @@ public class MongoDbFactoryBean implements FactoryBean<DB>, InitializingBean,
 		this.port = port;
 	}
 
-	public DB getObject() throws Exception {
+	public Mongo getObject() throws Exception {
 		Assert.notNull(mongo, "Mongo must not be null");
-		Assert.hasText(databaseName, "Database name must not be empty");
-		return mongo.getDB(databaseName);
+		return mongo;
 	}
 
-	public Class<? extends DB> getObjectType() {
-		return DB.class;
+	public Class<? extends Mongo> getObjectType() {
+		return Mongo.class;
 	}
 
 	public boolean isSingleton() {
@@ -80,17 +88,16 @@ public class MongoDbFactoryBean implements FactoryBean<DB>, InitializingBean,
 	public void afterPropertiesSet() throws Exception {
 		// apply defaults - convenient when used to configure for tests 
 		// in an application context
-		if (databaseName == null) {
-			logger.warn("Property databaseName not specified. Using default name 'test'");
-			databaseName = "test";
-		}
 		if (mongo == null) {
 			logger.warn("Property mongo not specified. Using default configuration");
 			if (host == null) {
 				mongo =  new Mongo();
 			}
 			else {
-				if (port == null) {
+				if(mongoOptions != null) {
+					mongo = new Mongo(host != null ? host : "localhost", mongoOptions);
+				}
+				else if (port == null) {
 					mongo = new Mongo(host);
 				}
 				else {

spring-datastore-mongodb/src/main/java/org/springframework/datastore/document/mongodb/MongoSynchronization.java

@@ -0,0 +1,15 @@
+package org.springframework.datastore.document.mongodb;
+
+import org.springframework.transaction.support.ResourceHolder;
+import org.springframework.transaction.support.ResourceHolderSynchronization;
+import org.springframework.transaction.support.TransactionSynchronization;
+
+public class MongoSynchronization extends ResourceHolderSynchronization
+		implements TransactionSynchronization {
+
+	public MongoSynchronization(ResourceHolder resourceHolder,
+			Object resourceKey) {
+		super(resourceHolder, resourceKey);
+	}
+
+}

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

@@ -31,6 +31,7 @@ import com.mongodb.CommandResult;
 import com.mongodb.DB;
 import com.mongodb.DBCollection;
 import com.mongodb.DBObject;
+import com.mongodb.Mongo;
 import com.mongodb.MongoException;
 import com.mongodb.WriteResult;
 import com.mongodb.util.JSON;
@@ -46,31 +47,38 @@ public class MongoTemplate extends AbstractDocumentStoreTemplate<DB> implements
 	//TODO expose configuration...
 	private CollectionOptions defaultCollectionOptions;
 
+	private Mongo mongo;
 
-	public MongoTemplate(DB db) {
-		super();
-		this.db = db;
+	private String databaseName;
+	
+	
+	public MongoTemplate(Mongo mongo, String databaseName) {
+		this(mongo, databaseName, null, null);
 	}
 	
-	public MongoTemplate(DB db, String defaultCollectionName) {
-		super();
-		this.db = db;
+	public MongoTemplate(Mongo mongo, String databaseName, String defaultCollectionName) {
+		this(mongo, databaseName, defaultCollectionName, null);
+	}
+	
+	public MongoTemplate(Mongo mongo, String databaseName, MongoConverter mongoConverter) {
+		this(mongo, databaseName, null, mongoConverter);
+	}
+	
+	public MongoTemplate(Mongo mongo, String databaseName, String defaultCollectionName, MongoConverter mongoConverter) {
+		this.mongoConverter = mongoConverter;
+		this.defaultCollectionName = defaultCollectionName;
+		this.mongo = mongo;
+		this.databaseName = databaseName;
+	}
+	
+	public void setDefaultCollectionName(String defaultCollectionName) {
 		this.defaultCollectionName = defaultCollectionName;
 	}
 
-	public MongoTemplate(DB db, MongoConverter mongoConverter) {
-		this(db);
-		this.mongoConverter = mongoConverter;
+	public void setDatabaseName(String databaseName) {
+		this.databaseName = databaseName;
 	}
 
-	public MongoTemplate(DB db, String defaultCollectionName, MongoConverter mongoConverter) {
-		this(db);
-		this.mongoConverter = mongoConverter;
-		this.defaultCollectionName = defaultCollectionName;
-	}
-	
-
-
 	public String getDefaultCollectionName() {
 		return defaultCollectionName;
 	}
@@ -268,7 +276,7 @@ public class MongoTemplate extends AbstractDocumentStoreTemplate<DB> implements
 
 	@Override
 	public DB getConnection() {
-		return db;
+		return MongoDbUtils.getDB(mongo, databaseName);
 	}
 
 	
@@ -292,6 +300,7 @@ public class MongoTemplate extends AbstractDocumentStoreTemplate<DB> implements
 
 	public void afterPropertiesSet() throws Exception {
 		if (this.getDefaultCollectionName() != null) {
+			DB db = getConnection();
 			if (! db.collectionExists(getDefaultCollectionName())) {
 				db.createCollection(getDefaultCollectionName(), null);
 			}

spring-datastore-mongodb/src/test/java/org/springframework/datastore/document/mongodb/analytics/MvcAnalyticsTests.java

@@ -34,8 +34,7 @@ public class MvcAnalyticsTests {
 	@Before
 	public void setUp() throws Exception {
 		Mongo m = new Mongo();
-		DB db = m.getDB("mvc");
-		mongoTemplate = new MongoTemplate(db, "mvc");
+		mongoTemplate = new MongoTemplate(m, "mvc", "mvc");
 		mongoTemplate.afterPropertiesSet();
 		
 		

spring-datastore-mongodb/template.mf

@@ -9,6 +9,7 @@ Import-Template:
  org.springframework.core.*;version="[3.0.0, 4.0.0)",
  org.springframework.dao.*;version="[3.0.0, 4.0.0)",
  org.springframework.util.*;version="[3.0.0, 4.0.0)",
+ org.springframework.transaction.*;version="[3.0.0, 4.0.0)",
  org.springframework.data.core.*;version="[1.0.0, 2.0.0)",
  org.springframework.datastore.core.*;version="[1.0.0, 2.0.0)",
  org.springframework.datastore.persistence.*;version="[1.0.0, 2.0.0)",