Generate and convert id on insert if explicitly defined.

We now make sure to provide an id value that matches the desired target type when no id is set, and the property defines an explicit conversion target.
Previously a new ObjectId would have been generated which leads to type inconsistencies when querying for _id.

Closes #4026
Original pull request: #4057.

Author: christophstrobl

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

@@ -97,6 +97,16 @@ public Document getDocument() {
 		return this.document;
 	}
 
+	/**
+	 * Updates the documents {@link #ID_FIELD}.
+	 *
+	 * @param value the {@literal _id} value to set.
+	 * @since 3.4.3
+	 */
+	public void updateId(Object value) {
+		document.put(ID_FIELD, value);
+	}
+
 	/**
 	 * An {@link UpdateDefinition} that indicates that the {@link #getUpdateObject() update object} has already been
 	 * mapped to the specific domain type.

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

@@ -1443,18 +1443,21 @@ protected Object insertDocument(String collectionName, Document document, Class<
 					collectionName));
 		}
 
+		MappedDocument mappedDocument = queryOperations.createInsertContext(MappedDocument.of(document))
+				.prepareId(entityClass);
+
 		return execute(collectionName, collection -> {
 			MongoAction mongoAction = new MongoAction(writeConcern, MongoActionOperation.INSERT, collectionName, entityClass,
-					document, null);
+					mappedDocument.getDocument(), null);
 			WriteConcern writeConcernToUse = prepareWriteConcern(mongoAction);
 
 			if (writeConcernToUse == null) {
-				collection.insertOne(document);
+				collection.insertOne(mappedDocument.getDocument());
 			} else {
-				collection.withWriteConcern(writeConcernToUse).insertOne(document);
+				collection.withWriteConcern(writeConcernToUse).insertOne(mappedDocument.getDocument());
 			}
 
-			return operations.forEntity(document).getId();
+			return operations.forEntity(mappedDocument.getDocument()).getId();
 		});
 	}
 
@@ -1505,7 +1508,9 @@ protected Object saveDocument(String collectionName, Document dbDoc, Class<?> en
 					: collection.withWriteConcern(writeConcernToUse);
 
 			if (!mapped.hasId()) {
-				collectionToUse.insertOne(dbDoc);
+
+				mapped = queryOperations.createInsertContext(mapped).prepareId(mappingContext.getPersistentEntity(entityClass));
+				collectionToUse.insertOne(mapped.getDocument());
 			} else {
 
 				MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);

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

@@ -28,7 +28,7 @@
 import org.bson.BsonValue;
 import org.bson.Document;
 import org.bson.codecs.Codec;
-
+import org.bson.types.ObjectId;
 import org.springframework.data.mapping.PropertyPath;
 import org.springframework.data.mapping.PropertyReferenceException;
 import org.springframework.data.mapping.context.MappingContext;
@@ -46,6 +46,7 @@
 import org.springframework.data.mongodb.core.aggregation.TypedAggregation;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.convert.UpdateMapper;
+import org.springframework.data.mongodb.core.mapping.MongoId;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
 import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 import org.springframework.data.mongodb.core.mapping.ShardKey;
@@ -107,6 +108,14 @@ class QueryOperations {
 		this.aggregationUtil = new AggregationUtil(queryMapper, mappingContext);
 	}
 
+	InsertContext createInsertContext(Document source) {
+		return createInsertContext(MappedDocument.of(source));
+	}
+
+	InsertContext createInsertContext(MappedDocument mappedDocument) {
+		return new InsertContext(mappedDocument);
+	}
+
 	/**
 	 * Create a new {@link QueryContext} instance.
 	 *
@@ -227,6 +236,57 @@ AggregationDefinition createAggregation(Aggregation aggregation,
 		return new AggregationDefinition(aggregation, aggregationOperationContext);
 	}
 
+	/**
+	 * {@link InsertContext} encapsulates common tasks required to interact with {@link Document} to be inserted.
+	 *
+	 * @since 3.4.3
+	 */
+	class InsertContext {
+
+		private final MappedDocument source;
+
+		private InsertContext(MappedDocument source) {
+			this.source = source;
+		}
+
+		/**
+		 * Prepare the {@literal _id} field. May generate a new {@literal id} value and convert it to the id properties
+		 * {@link MongoPersistentProperty#getFieldType() target type}.
+		 *
+		 * @param type must not be {@literal null}.
+		 * @param <T>
+		 * @return the {@link MappedDocument} containing the changes.
+		 * @see #prepareId(MongoPersistentEntity)
+		 */
+		<T> MappedDocument prepareId(Class<T> type) {
+			return prepareId(mappingContext.getPersistentEntity(type));
+		}
+
+		/**
+		 * Prepare the {@literal _id} field. May generate a new {@literal id} value and convert it to the id properties
+		 * {@link MongoPersistentProperty#getFieldType() target type}.
+		 *
+		 * @param entity can be {@literal null}.
+		 * @param <T>
+		 * @return the {@link MappedDocument} containing the changes.
+		 */
+		<T> MappedDocument prepareId(@Nullable MongoPersistentEntity<T> entity) {
+
+			if (entity == null) {
+				return source;
+			}
+
+			MongoPersistentProperty idProperty = entity.getIdProperty();
+			if (idProperty != null
+					&& (idProperty.hasExplicitWriteTarget() || idProperty.isAnnotationPresent(MongoId.class))) {
+				if (!ClassUtils.isAssignable(ObjectId.class, idProperty.getFieldType())) {
+					source.updateId(queryMapper.convertId(new ObjectId(), idProperty.getFieldType()));
+				}
+			}
+			return source;
+		}
+	}
+
 	/**
 	 * {@link QueryContext} encapsulates common tasks required to convert a {@link Query} into its MongoDB document
 	 * representation, mapping field names, as well as determining and applying {@link Collation collations}.
@@ -288,8 +348,7 @@ <T> Document getMappedQuery(@Nullable MongoPersistentEntity<T> entity) {
 			return queryMapper.getMappedObject(getQueryObject(), entity);
 		}
 
-		Document getMappedFields(@Nullable MongoPersistentEntity<?> entity,
-				EntityProjection<?, ?> projection) {
+		Document getMappedFields(@Nullable MongoPersistentEntity<?> entity, EntityProjection<?, ?> projection) {
 
 			Document fields = evaluateFields(entity);
 
@@ -402,8 +461,7 @@ private DistinctQueryContext(@Nullable Object query, String fieldName) {
 		}
 
 		@Override
-		Document getMappedFields(@Nullable MongoPersistentEntity<?> entity,
-				EntityProjection<?, ?> projection) {
+		Document getMappedFields(@Nullable MongoPersistentEntity<?> entity, EntityProjection<?, ?> projection) {
 			return getMappedFields(entity);
 		}
 

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

@@ -1448,7 +1448,8 @@ protected Mono<Object> insertDocument(String collectionName, Document dbDoc, Cla
 					.format("Inserting Document containing fields: " + dbDoc.keySet() + " in collection: " + collectionName));
 		}
 
-		Document document = new Document(dbDoc);
+		MappedDocument document = MappedDocument.of(dbDoc);
+		queryOperations.createInsertContext(document).prepareId(entityClass);
 
 		Flux<InsertOneResult> execute = execute(collectionName, collection -> {
 
@@ -1458,10 +1459,10 @@ protected Mono<Object> insertDocument(String collectionName, Document dbDoc, Cla
 
 			MongoCollection<Document> collectionToUse = prepareCollection(collection, writeConcernToUse);
 
-			return collectionToUse.insertOne(document);
+			return collectionToUse.insertOne(document.getDocument());
 		});
 
-		return Flux.from(execute).last().map(success -> MappedDocument.of(document).getId());
+		return Flux.from(execute).last().map(success -> document.getId());
 	}
 
 	protected Flux<ObjectId> insertDocumentList(String collectionName, List<Document> dbDocList) {
@@ -1525,7 +1526,7 @@ protected Mono<Object> saveDocument(String collectionName, Document document, Cl
 
 			Publisher<?> publisher;
 			if (!mapped.hasId()) {
-				publisher = collectionToUse.insertOne(document);
+				publisher = collectionToUse.insertOne(queryOperations.createInsertContext(mapped).prepareId(entityClass).getDocument());
 			} else {
 
 				MongoPersistentEntity<?> entity = mappingContext.getPersistentEntity(entityClass);

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

@@ -42,6 +42,7 @@
 import java.util.stream.IntStream;
 import java.util.stream.Stream;
 
+import org.bson.Document;
 import org.bson.types.ObjectId;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
@@ -3636,6 +3637,38 @@ public void saveAndLoadStringThatIsAnObjectIdAsString() {
 		assertThat(target).isEqualTo(source);
 	}
 
+	@Test // GH-4026
+	void saveShouldGenerateNewIdOfTypeIfExplicitlyDefined() {
+
+		RawStringId source = new RawStringId();
+		source.value = "new value";
+
+		template.save(source);
+
+		template.execute(RawStringId.class, collection -> {
+
+			org.bson.Document first = collection.find(new org.bson.Document()).first();
+			assertThat(first.get("_id")).isInstanceOf(String.class);
+			return null;
+		});
+	}
+
+	@Test // GH-4026
+	void insertShouldGenerateNewIdOfTypeIfExplicitlyDefined() {
+
+		RawStringId source = new RawStringId();
+		source.value = "new value";
+
+		template.insert(source);
+
+		template.execute(RawStringId.class, collection -> {
+
+			org.bson.Document first = collection.find(new org.bson.Document()).first();
+			assertThat(first.get("_id")).isInstanceOf(String.class);
+			return null;
+		});
+	}
+
 	@Test // DATAMONGO-2193
 	public void shouldNotConvertStringToObjectIdForNonIdField() {
 

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

@@ -18,6 +18,8 @@
 import static org.assertj.core.api.Assertions.*;
 import static org.mockito.Mockito.*;
 
+import org.bson.Document;
+import org.bson.types.ObjectId;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -32,7 +34,10 @@
 import org.springframework.data.mongodb.core.aggregation.TypeBasedAggregationOperationContext;
 import org.springframework.data.mongodb.core.convert.QueryMapper;
 import org.springframework.data.mongodb.core.convert.UpdateMapper;
+import org.springframework.data.mongodb.core.mapping.MongoId;
 import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentEntity;
+import org.springframework.data.mongodb.core.mapping.MongoPersistentProperty;
 
 /**
  * Unit tests for {@link QueryOperations}.
@@ -109,7 +114,99 @@ void createAggregationContextUsesStrictlyTypedContextForTypedAggregationsWhenReq
 		assertThat(ctx.getAggregationOperationContext()).isInstanceOf(TypeBasedAggregationOperationContext.class);
 	}
 
+	@Test // GH-4026
+	void insertContextDoesNotAddIdIfNoPersistentEntityCanBeFound() {
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one"));
+				});
+	}
+
+	@Test // GH-4026
+	void insertContextDoesNotAddIdIfNoIdPropertyCanBeFound() {
+
+		MongoPersistentEntity<Person> entity = mock(MongoPersistentEntity.class);
+		when(entity.getIdProperty()).thenReturn(null);
+		when(mappingContext.getPersistentEntity(eq(Person.class))).thenReturn((MongoPersistentEntity) entity);
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one"));
+				});
+	}
+
+	@Test // GH-4026
+	void insertContextDoesNotAddConvertedIdForNonExplicitFieldTypes() {
+
+		MongoPersistentEntity<Person> entity = mock(MongoPersistentEntity.class);
+		MongoPersistentProperty property = mock(MongoPersistentProperty.class);
+		when(entity.getIdProperty()).thenReturn(property);
+		when(property.hasExplicitWriteTarget()).thenReturn(false);
+		doReturn(entity).when(mappingContext).getPersistentEntity(eq(Person.class));
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one"));
+				});
+	}
+
+	@Test // GH-4026
+	void insertContextAddsConvertedIdForExplicitFieldTypes() {
+
+		MongoPersistentEntity<Person> entity = mock(MongoPersistentEntity.class);
+		MongoPersistentProperty property = mock(MongoPersistentProperty.class);
+		when(entity.getIdProperty()).thenReturn(property);
+		when(property.hasExplicitWriteTarget()).thenReturn(true);
+		doReturn(String.class).when(property).getFieldType();
+		doReturn(entity).when(mappingContext).getPersistentEntity(eq(Person.class));
+
+		when(queryMapper.convertId(any(), eq(String.class))).thenReturn("&#9774;");
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one").append("_id", "&#9774;"));
+				});
+	}
+
+	@Test // GH-4026
+	void insertContextAddsConvertedIdForMongoIdTypes() {
+
+		MongoPersistentEntity<Person> entity = mock(MongoPersistentEntity.class);
+		MongoPersistentProperty property = mock(MongoPersistentProperty.class);
+		when(entity.getIdProperty()).thenReturn(property);
+		when(property.hasExplicitWriteTarget()).thenReturn(false);
+		when(property.isAnnotationPresent(eq(MongoId.class))).thenReturn(true);
+		doReturn(String.class).when(property).getFieldType();
+		doReturn(entity).when(mappingContext).getPersistentEntity(eq(Person.class));
+
+		when(queryMapper.convertId(any(), eq(String.class))).thenReturn("&#9774;");
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one").append("_id", "&#9774;"));
+				});
+	}
+
+	@Test // GH-4026
+	void insertContextDoesNotAddConvertedIdForMongoIdTypesTargetingObjectId() {
+
+		MongoPersistentEntity<Person> entity = mock(MongoPersistentEntity.class);
+		MongoPersistentProperty property = mock(MongoPersistentProperty.class);
+		when(entity.getIdProperty()).thenReturn(property);
+		when(property.hasExplicitWriteTarget()).thenReturn(false);
+		when(property.isAnnotationPresent(eq(MongoId.class))).thenReturn(true);
+		doReturn(ObjectId.class).when(property).getFieldType();
+		doReturn(entity).when(mappingContext).getPersistentEntity(eq(Person.class));
+
+		assertThat(queryOperations.createInsertContext(new Document("value", "one")).prepareId(Person.class).getDocument())//
+				.satisfies(result -> {
+					assertThat(result).isEqualTo(new Document("value", "one"));
+				});
+	}
+
 	static class Person {
 
 	}
+
 }