Support generating JsonSchema for Polymorphic fields.

This commit introduces MergedJsonSchema and MergedJsonSchemaProperty that can be used to merge properties of multiple objects into one as long as the additions do not conflict with another (eg. due to usage of different types).
To resolve previously mentioned errors it is required to provide a ConflictResolutionFunction.

Closes #3870
Original pull request: #3986.

Author: christophstrobl

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

@@ -24,7 +24,6 @@
 import java.util.stream.Collectors;
 
 import org.bson.Document;
-
 import org.springframework.data.mapping.PersistentProperty;
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.mongodb.core.convert.MongoConverter;
@@ -45,6 +44,7 @@
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.CollectionUtils;
+import org.springframework.util.LinkedMultiValueMap;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.StringUtils;
 
@@ -62,6 +62,7 @@ class MappingMongoJsonSchemaCreator implements MongoJsonSchemaCreator {
 	private final MongoConverter converter;
 	private final MappingContext<MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext;
 	private final Predicate<JsonSchemaPropertyContext> filter;
+	private final LinkedMultiValueMap<String, Class<?>> mergeProperties;
 
 	/**
 	 * Create a new instance of {@link MappingMongoJsonSchemaCreator}.
@@ -72,23 +73,51 @@ class MappingMongoJsonSchemaCreator implements MongoJsonSchemaCreator {
 	MappingMongoJsonSchemaCreator(MongoConverter converter) {
 
 		this(converter, (MappingContext<MongoPersistentEntity<?>, MongoPersistentProperty>) converter.getMappingContext(),
-				(property) -> true);
+				(property) -> true, new LinkedMultiValueMap<>());
 	}
 
 	@SuppressWarnings("unchecked")
 	MappingMongoJsonSchemaCreator(MongoConverter converter,
 			MappingContext<MongoPersistentEntity<?>, MongoPersistentProperty> mappingContext,
-			Predicate<JsonSchemaPropertyContext> filter) {
+			Predicate<JsonSchemaPropertyContext> filter, LinkedMultiValueMap<String, Class<?>> mergeProperties) {
 
 		Assert.notNull(converter, "Converter must not be null!");
 		this.converter = converter;
 		this.mappingContext = mappingContext;
 		this.filter = filter;
+		this.mergeProperties = mergeProperties;
 	}
 
 	@Override
 	public MongoJsonSchemaCreator filter(Predicate<JsonSchemaPropertyContext> filter) {
-		return new MappingMongoJsonSchemaCreator(converter, mappingContext, filter);
+		return new MappingMongoJsonSchemaCreator(converter, mappingContext, filter, mergeProperties);
+	}
+
+	@Override
+	public PropertySpecifier specify(String path) {
+		return new PropertySpecifier() {
+			@Override
+			public MongoJsonSchemaCreator types(Class<?>... types) {
+				return specifyTypesFor(path, types);
+			}
+		};
+	}
+
+	/**
+	 * Specify additional types to be considered wehen rendering the schema for the given path.
+	 *
+	 * @param path path the path using {@literal dot '.'} notation.
+	 * @param types must not be {@literal null}.
+	 * @return new instance of {@link MongoJsonSchemaCreator}.
+	 * @since 3.4
+	 */
+	public MongoJsonSchemaCreator specifyTypesFor(String path, Class<?>... types) {
+
+		LinkedMultiValueMap<String, Class<?>> clone = mergeProperties.clone();
+		for (Class<?> type : types) {
+			clone.add(path, type);
+		}
+		return new MappingMongoJsonSchemaCreator(converter, mappingContext, filter, clone);
 	}
 
 	@Override
@@ -131,9 +160,12 @@ private List<JsonSchemaProperty> computePropertiesForEntity(List<MongoPersistent
 
 			List<MongoPersistentProperty> currentPath = new ArrayList<>(path);
 
-			if (!filter.test(new PropertyContext(
-					currentPath.stream().map(PersistentProperty::getName).collect(Collectors.joining(".")), nested))) {
-				continue;
+			String stringPath = currentPath.stream().map(PersistentProperty::getName).collect(Collectors.joining("."));
+			stringPath = StringUtils.hasText(stringPath) ? (stringPath + "." + nested.getName()) : nested.getName();
+			if (!filter.test(new PropertyContext(stringPath, nested))) {
+				if (!mergeProperties.containsKey(stringPath)) {
+					continue;
+				}
 			}
 
 			if (path.contains(nested)) { // cycle guard
@@ -151,14 +183,34 @@ private List<JsonSchemaProperty> computePropertiesForEntity(List<MongoPersistent
 
 	private JsonSchemaProperty computeSchemaForProperty(List<MongoPersistentProperty> path) {
 
+		String stringPath = path.stream().map(MongoPersistentProperty::getName).collect(Collectors.joining("."));
 		MongoPersistentProperty property = CollectionUtils.lastElement(path);
 
 		boolean required = isRequiredProperty(property);
 		Class<?> rawTargetType = computeTargetType(property); // target type before conversion
 		Class<?> targetType = converter.getTypeMapper().getWriteTargetTypeFor(rawTargetType); // conversion target type
 
-		if (!isCollection(property) && property.isEntity() && ObjectUtils.nullSafeEquals(rawTargetType, targetType)) {
-			return createObjectSchemaPropertyForEntity(path, property, required);
+		if (!isCollection(property) && ObjectUtils.nullSafeEquals(rawTargetType, targetType)) {
+			if (property.isEntity() || mergeProperties.containsKey(stringPath)) {
+				List<JsonSchemaProperty> targetProperties = new ArrayList<>();
+
+				if (property.isEntity()) {
+					targetProperties.add(createObjectSchemaPropertyForEntity(path, property, required));
+				}
+				if (mergeProperties.containsKey(stringPath)) {
+					for (Class<?> theType : mergeProperties.get(stringPath)) {
+
+						ObjectJsonSchemaProperty target = JsonSchemaProperty.object(property.getName());
+						List<JsonSchemaProperty> nestedProperties = computePropertiesForEntity(path,
+								mappingContext.getRequiredPersistentEntity(theType));
+
+						targetProperties.add(createPotentiallyRequiredSchemaProperty(
+								target.properties(nestedProperties.toArray(new JsonSchemaProperty[0])), required));
+					}
+				}
+				return targetProperties.size() == 1 ? targetProperties.iterator().next()
+						: JsonSchemaProperty.combined(targetProperties);
+			}
 		}
 
 		String fieldName = computePropertyFieldName(property);

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

@@ -15,6 +15,7 @@
  */
 package org.springframework.data.mongodb.core;
 
+import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.function.Predicate;
@@ -62,7 +63,6 @@
  * {@link org.springframework.data.annotation.Id _id} properties using types that can be converted into
  * {@link org.bson.types.ObjectId} like {@link String} will be mapped to {@code type : 'object'} unless there is more
  * specific information available via the {@link org.springframework.data.mongodb.core.mapping.MongoId} annotation.
-
  * {@link Encrypted} properties will contain {@literal encrypt} information.
  *
  * @author Christoph Strobl
@@ -78,6 +78,20 @@ public interface MongoJsonSchemaCreator {
 	 */
 	MongoJsonSchema createSchemaFor(Class<?> type);
 
+	/**
+	 * Create a combined {@link MongoJsonSchema} out of the individual schemas of the given types by combining their
+	 * properties into one large {@link MongoJsonSchema schema}.
+	 * 
+	 * @param types must not be {@literal null} nor contain {@literal null}.
+	 * @return new instance of {@link MongoJsonSchema}.
+	 * @since 3.4
+	 */
+	default MongoJsonSchema combineSchemaFor(Class<?>... types) {
+
+		MongoJsonSchema[] schemas = Arrays.stream(types).map(this::createSchemaFor).toArray(MongoJsonSchema[]::new);
+		return MongoJsonSchema.combined(schemas);
+	}
+
 	/**
 	 * Filter matching {@link JsonSchemaProperty properties}.
 	 *
@@ -87,6 +101,15 @@ public interface MongoJsonSchemaCreator {
 	 */
 	MongoJsonSchemaCreator filter(Predicate<JsonSchemaPropertyContext> filter);
 
+	/**
+	 * Entry point to specify additional behavior for a given path.
+	 *
+	 * @param path the path using {@literal dot '.'} notation.
+	 * @return new instance of {@link PropertySpecifier}.
+	 * @since 3.4
+	 */
+	PropertySpecifier specify(String path);
+
 	/**
 	 * The context in which a specific {@link #getProperty()} is encountered during schema creation.
 	 * 
@@ -209,4 +232,20 @@ static MongoJsonSchemaCreator create() {
 
 		return create(converter);
 	}
+
+	/**
+	 * @since 3.4
+	 * @author Christoph Strobl
+	 * @since 3.4
+	 */
+	interface PropertySpecifier {
+
+		/**
+		 * Set additional type parameters for polymorphic ones.
+		 *
+		 * @param types must not be {@literal null}.
+		 * @return the source
+		 */
+		MongoJsonSchemaCreator types(Class<?>... types);
+	}
 }

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.schema;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BiFunction;
+
+import org.bson.Document;
+
+/**
+ * {@link MongoJsonSchema} implementation that is capable of combining properties of different schemas into one.
+ *
+ * @author Christoph Strobl
+ * @since 3.4
+ */
+class CombinedJsonSchema implements MongoJsonSchema {
+
+	private final List<MongoJsonSchema> schemaList;
+	private final BiFunction<Map<String, Object>, Map<String, Object>, Document> mergeFunction;
+
+	CombinedJsonSchema(List<MongoJsonSchema> schemaList, ConflictResolutionFunction conflictResolutionFunction) {
+		this(schemaList, new TypeUnifyingMergeFunction(conflictResolutionFunction));
+	}
+
+	CombinedJsonSchema(List<MongoJsonSchema> schemaList,
+			BiFunction<Map<String, Object>, Map<String, Object>, Document> mergeFunction) {
+
+		this.schemaList = new ArrayList<>(schemaList);
+		this.mergeFunction = mergeFunction;
+	}
+
+	@Override
+	public MongoJsonSchema combineWith(Collection<MongoJsonSchema> sources) {
+
+		schemaList.addAll(sources);
+		return this;
+	}
+
+	@Override
+	public Document schemaDocument() {
+
+		Document targetSchema = new Document();
+		for (MongoJsonSchema schema : schemaList) {
+			targetSchema = mergeFunction.apply(targetSchema, schema.schemaDocument());
+		}
+
+		return targetSchema;
+	}
+}

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

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.mongodb.core.schema;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.BiFunction;
+
+import org.bson.Document;
+import org.springframework.data.mongodb.core.schema.MongoJsonSchema.ConflictResolutionFunction;
+
+/**
+ * {@link JsonSchemaProperty} implementation that is capable of combining multiple properties with different values into
+ * a single one.
+ *
+ * @author Christoph Strobl
+ * @since 3.4
+ */
+class CombinedJsonSchemaProperty implements JsonSchemaProperty {
+
+	private final Iterable<JsonSchemaProperty> properties;
+	private final BiFunction<Map<String, Object>, Map<String, Object>, Document> mergeFunction;
+
+	CombinedJsonSchemaProperty(Iterable<JsonSchemaProperty> properties) {
+		this(properties, (k, a, b) -> {
+			throw new IllegalStateException(
+					String.format("Error resolving conflict for %s. No conflict resolution function defined.", k));
+		});
+	}
+
+	CombinedJsonSchemaProperty(Iterable<JsonSchemaProperty> properties,
+			ConflictResolutionFunction conflictResolutionFunction) {
+		this(properties, new TypeUnifyingMergeFunction(conflictResolutionFunction));
+	}
+
+	CombinedJsonSchemaProperty(Iterable<JsonSchemaProperty> properties,
+			BiFunction<Map<String, Object>, Map<String, Object>, Document> mergeFunction) {
+
+		this.properties = properties;
+		this.mergeFunction = mergeFunction;
+	}
+
+	@Override
+	public Set<Type> getTypes() {
+		return Collections.emptySet();
+	}
+
+	@Override
+	public Document toDocument() {
+
+		Document document = new Document();
+
+		for (JsonSchemaProperty property : properties) {
+			document = mergeFunction.apply(document, property.toDocument());
+		}
+		return document;
+	}
+
+	@Override
+	public String getIdentifier() {
+		return properties.iterator().next().getIdentifier();
+	}
+}

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

@@ -15,6 +15,8 @@
  */
 package org.springframework.data.mongodb.core.schema;
 
+import java.util.Collection;
+
 import org.springframework.data.mongodb.core.schema.TypedJsonSchemaObject.NumericJsonSchemaObject;
 import org.springframework.data.mongodb.core.schema.TypedJsonSchemaObject.ObjectJsonSchemaObject;
 import org.springframework.data.mongodb.core.schema.IdentifiableJsonSchemaProperty.*;
@@ -233,6 +235,17 @@ static JsonSchemaProperty required(JsonSchemaProperty property) {
 		return new RequiredJsonSchemaProperty(property, true);
 	}
 
+	/**
+	 * Combines multiple {@link JsonSchemaProperty} with potentially different attributes into one.
+	 *
+	 * @param properties must not be {@literal null}.
+	 * @return new instance of {@link JsonSchemaProperty}.
+	 * @since 3.4
+	 */
+	static JsonSchemaProperty combined(Collection<JsonSchemaProperty> properties) {
+		return new CombinedJsonSchemaProperty(properties);
+	}
+
 	/**
 	 * Builder for {@link IdentifiableJsonSchemaProperty}.
 	 */