DATAMONGO-1513 - Fixed identifier population for event listener generated, non-ObjectId on batch inserts.

The methods in MongoTemplate inserting a batch of documents previously only returned database generated identifiers, more especially ObjectId ones. This caused non-ObjectId identifiers potentially generated by other parties — i.e. an event listener reacting to a BeforeSaveEvent — not being considered for source object identifier population.

This commit adds a workaround augmenting the list of database generated identifiers with the ones actually present in the documents to be inserted. A follow-up ticket DATAMONGO-1519 was created to track the removal of the workaround in favor of a proper fix unfortunately requiring a change in public API (so a 2.0 candidate only).

Related tickets: DATAMONGO-1519.

Author: odrotbohm

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

@@ -948,7 +948,9 @@ protected <T> void doInsertBatch(String collectionName, Collection<? extends T>
 			maybeEmitEvent(new BeforeSaveEvent<T>(o, dbDoc, collectionName));
 			dbObjectList.add(dbDoc);
 		}
-		List<ObjectId> ids = insertDBObjectList(collectionName, dbObjectList);
+
+		List<Object> ids = consolidateIdentifiers(insertDBObjectList(collectionName, dbObjectList), dbObjectList);
+
 		int i = 0;
 		for (T obj : batchToSave) {
 			if (i < ids.size()) {
@@ -1053,6 +1055,8 @@ public Object doInCollection(DBCollection collection) throws MongoException, Dat
 		});
 	}
 
+	// TODO: 2.0 - Change method signature to return List<Object> and return all identifiers (DATAMONGO-1513,
+	// DATAMONGO-1519)
 	protected List<ObjectId> insertDBObjectList(final String collectionName, final List<DBObject> dbDocList) {
 		if (dbDocList.isEmpty()) {
 			return Collections.emptyList();
@@ -2131,6 +2135,28 @@ private static RuntimeException potentiallyConvertRuntimeException(RuntimeExcept
 		return resolved == null ? ex : resolved;
 	}
 
+	/**
+	 * Returns all identifiers for the given documents. Will augment the given identifiers and fill in only the ones that
+	 * are {@literal null} currently. This would've been better solved in {@link #insertDBObjectList(String, List)}
+	 * directly but would require a signature change of that method.
+	 * 
+	 * @param ids
+	 * @param documents
+	 * @return TODO: Remove for 2.0 and change method signature of {@link #insertDBObjectList(String, List)}.
+	 */
+	private static List<Object> consolidateIdentifiers(List<ObjectId> ids, List<DBObject> documents) {
+
+		List<Object> result = new ArrayList<Object>(ids.size());
+
+		for (int i = 0; i < ids.size(); i++) {
+
+			ObjectId objectId = ids.get(i);
+			result.add(objectId == null ? documents.get(i).get(ID_FIELD) : objectId);
+		}
+
+		return result;
+	}
+
 	// Callback implementations
 
 	/**

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

@@ -37,6 +37,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.UUID;
 
 import org.bson.types.ObjectId;
 import org.joda.time.DateTime;
@@ -47,6 +48,7 @@
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.DataIntegrityViolationException;
@@ -74,11 +76,14 @@
 import org.springframework.data.mongodb.core.index.IndexInfo;
 import org.springframework.data.mongodb.core.mapping.Field;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.data.mongodb.core.mapping.event.AbstractMongoEventListener;
+import org.springframework.data.mongodb.core.mapping.event.BeforeSaveEvent;
 import org.springframework.data.mongodb.core.query.BasicQuery;
 import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.data.util.CloseableIterator;
+import org.springframework.test.annotation.DirtiesContext;
 import org.springframework.test.context.ContextConfiguration;
 import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 import org.springframework.util.ObjectUtils;
@@ -123,6 +128,7 @@ public class MongoTemplateTests {
 
 	@Autowired MongoTemplate template;
 	@Autowired MongoDbFactory factory;
+	@Autowired ConfigurableApplicationContext context;
 
 	MongoTemplate mappingTemplate;
 	org.springframework.data.util.Version mongoVersion;
@@ -3442,6 +3448,28 @@ public void shouldFetchMapOfLazyReferencesCorrectly() {
 		assertThat(target.lazyDbRefAnnotatedMap.values(), contains(two, one));
 	}
 
+	/**
+	 * @see DATAMONGO-1513
+	 */
+	@Test
+	@DirtiesContext
+	public void populatesIdsAddedByEventListener() {
+
+		context.addApplicationListener(new AbstractMongoEventListener<Document>() {
+
+			@Override
+			public void onBeforeSave(BeforeSaveEvent<Document> event) {
+				event.getDBObject().put("_id", UUID.randomUUID().toString());
+			}
+		});
+
+		Document document = new Document();
+
+		template.insertAll(Arrays.asList(document));
+
+		assertThat(document.id, is(notNullValue()));
+	}
+
 	static class TypeWithNumbers {
 
 		@Id String id;