Allow keys other than _id for primary key in MongoItemWriter

When moving data between different data store types, such as
between MongoDB and a SQL database, not all systems have a
"_id" as their pimary key.  Most often it has a different
name or can be multiple keys/columns used as a composite
key.  This change allows the app to specify an alternate
key or a set of keys which, together, can be used as the
primary key for a MongoDB insert/update/upsert.

Signed-off-by: blake_bauman <[email protected]>

Author: blake-bauman

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/MongoItemWriter.java

@@ -16,6 +16,10 @@
 
 package org.springframework.batch.item.data;
 
+import static java.util.stream.Collectors.toList;
+
+import java.util.List;
+
 import org.bson.Document;
 import org.bson.types.ObjectId;
 
@@ -92,6 +96,8 @@ public enum Mode {
 
 	private Mode mode = Mode.UPSERT;
 
+	private List<String> primaryKeys = List.of(ID_KEY);
+
 	public MongoItemWriter() {
 		super();
 		this.bufferKey = new Object();
@@ -150,6 +156,27 @@ public String getCollection() {
 		return collection;
 	}
 
+	/**
+	 * Set the primary keys to associate with the document being written. These fields
+	 * should uniquely identify a single object.
+	 * @param primaryKeys The primary keys to use.
+	 * @since 5.2.3
+	 */
+	public void setPrimaryKeys(List<String> primaryKeys) {
+		Assert.notEmpty(primaryKeys, "The primaryKeys list must have one or more keys.");
+
+		this.primaryKeys = primaryKeys;
+	}
+
+	/**
+	 * Get the list of primary keys associated with the document being written.
+	 * @return the list of primary keys
+	 * @since 5.2.3
+	 */
+	public List<String> getPrimaryKeys() {
+		return primaryKeys;
+	}
+
 	/**
 	 * If a transaction is active, buffer items to be written just before commit.
 	 * Otherwise write items using the provided template.
@@ -200,9 +227,14 @@ private void remove(Chunk<? extends T> chunk) {
 		for (Object item : chunk) {
 			Document document = new Document();
 			mongoConverter.write(item, document);
-			Object objectId = document.get(ID_KEY);
-			if (objectId != null) {
-				Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId));
+
+			List<Criteria> criteriaList = primaryKeys.stream()
+				.filter(document::containsKey)
+				.map(key -> Criteria.where(key).is(document.get(key)))
+				.collect(toList());
+			if (!criteriaList.isEmpty()) {
+				Query query = new Query();
+				criteriaList.forEach(query::addCriteria);
 				bulkOperations.remove(query);
 			}
 		}
@@ -216,8 +248,21 @@ private void upsert(Chunk<? extends T> chunk) {
 		for (Object item : chunk) {
 			Document document = new Document();
 			mongoConverter.write(item, document);
-			Object objectId = document.get(ID_KEY) != null ? document.get(ID_KEY) : new ObjectId();
-			Query query = new Query().addCriteria(Criteria.where(ID_KEY).is(objectId));
+
+			Query query = new Query();
+			List<Criteria> criteriaList = primaryKeys.stream()
+				.filter(document::containsKey)
+				.map(key -> Criteria.where(key).is(document.get(key)))
+				.collect(toList());
+
+			if (criteriaList.isEmpty()) {
+				Object objectId = document.get(ID_KEY) != null ? document.get(ID_KEY) : new ObjectId();
+				query.addCriteria(Criteria.where(ID_KEY).is(objectId));
+			}
+			else {
+				criteriaList.forEach(query::addCriteria);
+			}
+
 			bulkOperations.replaceOne(query, document, upsert);
 		}
 		bulkOperations.execute();

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/data/builder/MongoItemWriterBuilder.java

@@ -16,6 +16,8 @@
 
 package org.springframework.batch.item.data.builder;
 
+import java.util.List;
+
 import org.springframework.batch.item.data.MongoItemWriter;
 import org.springframework.batch.item.data.MongoItemWriter.Mode;
 import org.springframework.data.mongodb.core.MongoOperations;
@@ -37,6 +39,8 @@ public class MongoItemWriterBuilder<T> {
 
 	private Mode mode = Mode.UPSERT;
 
+	private List<String> primaryKeys = List.of();
+
 	/**
 	 * Set the operating {@link Mode} to be applied by this writer. Defaults to
 	 * {@link Mode#UPSERT}.
@@ -76,6 +80,32 @@ public MongoItemWriterBuilder<T> collection(String collection) {
 		return this;
 	}
 
+	/**
+	 * Set the primary keys to associate with the document being written. These fields
+	 * should uniquely identify a single object.
+	 * @param primaryKeys The keys to use.
+	 * @see MongoItemWriter#setPrimaryKeys(List)
+	 * @since 5.2.3
+	 */
+	public MongoItemWriterBuilder<T> primaryKeys(List<String> primaryKeys) {
+		this.primaryKeys = List.copyOf(primaryKeys);
+
+		return this;
+	}
+
+	/**
+	 * Set the primary keys to associate with the document being written. These fields
+	 * should uniquely identify a single object.
+	 * @param primaryKeys The keys to use.
+	 * @see MongoItemWriter#setPrimaryKeys(List)
+	 * @since 5.2.3
+	 */
+	public MongoItemWriterBuilder<T> primaryKeys(String... primaryKeys) {
+		this.primaryKeys = List.of(primaryKeys);
+
+		return this;
+	}
+
 	/**
 	 * Validates and builds a {@link MongoItemWriter}.
 	 * @return a {@link MongoItemWriter}
@@ -88,6 +118,10 @@ public MongoItemWriter<T> build() {
 		writer.setMode(this.mode);
 		writer.setCollection(this.collection);
 
+		if (!this.primaryKeys.isEmpty()) {
+			writer.setPrimaryKeys(this.primaryKeys);
+		}
+
 		return writer;
 	}