HSEARCH-4153 Implement loading for the JavaBean mapper

Signed-off-by: Yoann Rodière <yoann@hibernate.org>

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/EntityLoader.java

@@ -0,0 +1,22 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.mapper.javabean.loading;
+
+import java.util.List;
+
+public interface EntityLoader<E> {
+
+	/**
+	 * Loads the entities corresponding to the given identifiers, blocking the current thread while doing so.
+	 *
+	 * @param identifiers A list of identifiers for objects to load.
+	 * @return A list of entities, in the same order the identifiers were given.
+	 * {@code null} is inserted when an object is not found or has the wrong concrete type.
+	 */
+	List<E> load(List<?> identifiers);
+
+}

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/LoadingOptions.java

@@ -0,0 +1,13 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.mapper.javabean.loading;
+
+public interface LoadingOptions {
+
+	<T> LoadingOptions registerLoader(Class<T> type, EntityLoader<T> loader);
+
+}

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/impl/JavaBeanLoader.java

@@ -0,0 +1,39 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.mapper.javabean.loading.impl;
+
+import java.util.List;
+
+import org.hibernate.search.engine.common.timing.spi.Deadline;
+import org.hibernate.search.mapper.javabean.loading.EntityLoader;
+import org.hibernate.search.mapper.pojo.loading.spi.PojoLoader;
+import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeIdentifier;
+
+public class JavaBeanLoader<E> implements PojoLoader<E> {
+
+	private final EntityLoader<E> delegate;
+
+	public JavaBeanLoader(EntityLoader<E> delegate) {
+		this.delegate = delegate;
+	}
+
+	@Override
+	public List<?> loadBlocking(List<?> identifiers, Deadline deadline) {
+		return delegate.load( identifiers );
+	}
+
+	@Override
+	@SuppressWarnings("unchecked")
+	public <E2 extends E> E2 castToExactTypeOrNull(PojoRawTypeIdentifier<E2> expectedType, Object loadedObject) {
+		if ( expectedType.javaClass().equals( loadedObject.getClass() ) ) {
+			return (E2) loadedObject;
+		}
+		else {
+			return null;
+		}
+	}
+}

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/impl/JavaBeanSearchLoadingContext.java

@@ -7,8 +7,13 @@
 package org.hibernate.search.mapper.javabean.loading.impl;
 
 import java.lang.invoke.MethodHandles;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.Set;
 
+import org.hibernate.search.mapper.javabean.loading.EntityLoader;
+import org.hibernate.search.mapper.javabean.loading.LoadingOptions;
 import org.hibernate.search.mapper.javabean.log.impl.Log;
 import org.hibernate.search.mapper.pojo.loading.spi.PojoLoader;
 import org.hibernate.search.mapper.pojo.loading.spi.PojoLoadingContext;
@@ -20,7 +25,10 @@ public final class JavaBeanSearchLoadingContext implements PojoLoadingContext {
 
 	private static final Log log = LoggerFactory.make( Log.class, MethodHandles.lookup() );
 
-	private JavaBeanSearchLoadingContext() {
+	private final Map<PojoRawTypeIdentifier<?>, PojoLoader<?>> loaderByType;
+
+	private JavaBeanSearchLoadingContext(Builder builder) {
+		this.loaderByType = builder.loaderByType == null ? Collections.emptyMap() : builder.loaderByType;
 	}
 
 	@Override
@@ -35,22 +43,45 @@ public Object loaderKey(PojoRawTypeIdentifier<?> type) {
 	}
 
 	@Override
+	@SuppressWarnings("unchecked")
 	public <T> PojoLoader<T> createLoader(Set<PojoRawTypeIdentifier<? extends T>> expectedTypes) {
-		throw log.entityLoadingNotSupported();
+		PojoRawTypeIdentifier<? extends T> type = expectedTypes.iterator().next();
+		PojoLoader<T> loader = (PojoLoader<T>) loaderByType.get( type );
+		if ( loader == null ) {
+			throw log.entityLoaderNotRegistered( type );
+		}
+		return loader;
 	}
 
-	public static final class Builder implements PojoLoadingContextBuilder<Void> {
-		public Builder() {
+	public static final class Builder implements PojoLoadingContextBuilder<LoadingOptions>, LoadingOptions {
+		private final LoadingTypeContextProvider typeContextProvider;
+		private Map<PojoRawTypeIdentifier<?>, PojoLoader<?>> loaderByType;
+
+		public Builder(LoadingTypeContextProvider typeContextProvider) {
+			this.typeContextProvider = typeContextProvider;
+		}
+
+		@Override
+		public LoadingOptions toAPI() {
+			return this;
 		}
 
 		@Override
-		public Void toAPI() {
-			throw log.entityLoadingNotSupported();
+		public <T> LoadingOptions registerLoader(Class<T> type, EntityLoader<T> loader) {
+			LoadingTypeContext<T> typeContext = typeContextProvider.indexedForExactClass( type );
+			if ( typeContext == null ) {
+				throw log.notIndexedEntityType( type );
+			}
+			if ( loaderByType == null ) {
+				loaderByType = new LinkedHashMap<>();
+			}
+			loaderByType.put( typeContext.typeIdentifier(), new JavaBeanLoader<>( loader ) );
+			return this;
 		}
 
 		@Override
 		public PojoLoadingContext build() {
-			return new JavaBeanSearchLoadingContext();
+			return new JavaBeanSearchLoadingContext( this );
 		}
 	}
 }

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/impl/LoadingTypeContext.java

@@ -0,0 +1,15 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.mapper.javabean.loading.impl;
+
+import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeIdentifier;
+
+public interface LoadingTypeContext<E> {
+
+	PojoRawTypeIdentifier<E> typeIdentifier();
+
+}

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/loading/impl/LoadingTypeContextProvider.java

@@ -0,0 +1,13 @@
+/*
+ * Hibernate Search, full-text search for your domain model
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+package org.hibernate.search.mapper.javabean.loading.impl;
+
+public interface LoadingTypeContextProvider {
+
+	<E> LoadingTypeContext<E> indexedForExactClass(Class<E> typeIdentifier);
+
+}

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/log/impl/Log.java

@@ -9,6 +9,7 @@
 
 import org.hibernate.search.engine.backend.common.DocumentReference;
 import org.hibernate.search.engine.environment.bean.spi.BeanProvider;
+import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeIdentifier;
 import org.hibernate.search.util.common.SearchException;
 import org.hibernate.search.util.common.logging.impl.MessageConstants;
 import org.hibernate.search.util.common.logging.impl.ClassFormatter;
@@ -45,8 +46,8 @@ public interface Log extends BasicLogger {
 	SearchException namedTypesNotSupported(String name);
 
 	@Message(id = ID_OFFSET + 8,
-			value = "Unable to set up entity loading: the JavaBean mapper does not support entity loading.")
-	SearchException entityLoadingNotSupported();
+			value = "Unable to set up entity loading for type '%s', because no entity loader was registered for this type.")
+	SearchException entityLoaderNotRegistered(PojoRawTypeIdentifier<?> typeIdentifier);
 
 	@Message(id = ID_OFFSET + 9, value = "Type '%1$s' is not an entity type, or this entity type is not indexed.")
 	SearchException notIndexedEntityType(@FormatWith(ClassFormatter.class) Class<?> type);

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/mapping/impl/JavaBeanIndexedTypeContext.java

@@ -8,6 +8,7 @@
 
 import org.hibernate.search.engine.backend.index.IndexManager;
 import org.hibernate.search.engine.mapper.mapping.spi.MappedIndexManager;
+import org.hibernate.search.mapper.javabean.loading.impl.LoadingTypeContext;
 import org.hibernate.search.mapper.javabean.scope.impl.JavaBeanScopeIndexedTypeContext;
 import org.hibernate.search.mapper.javabean.session.impl.JavaBeanSessionIndexedTypeContext;
 import org.hibernate.search.mapper.pojo.bridge.runtime.spi.IdentifierMapping;
@@ -16,7 +17,7 @@
 import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeIdentifier;
 
 class JavaBeanIndexedTypeContext<E> extends AbstractJavaBeanTypeContext<E>
-		implements JavaBeanScopeIndexedTypeContext<E>, JavaBeanSessionIndexedTypeContext<E> {
+		implements JavaBeanScopeIndexedTypeContext<E>, JavaBeanSessionIndexedTypeContext<E>, LoadingTypeContext<E> {
 	private final IdentifierMapping identifierMapping;
 	private final MappedIndexManager indexManager;
 

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/mapping/impl/JavaBeanTypeContextContainer.java

@@ -12,11 +12,12 @@
 import java.util.List;
 import java.util.Map;
 
+import org.hibernate.search.mapper.javabean.loading.impl.LoadingTypeContextProvider;
 import org.hibernate.search.mapper.javabean.session.impl.JavaBeanSearchSessionTypeContextProvider;
 import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeIdentifier;
 import org.hibernate.search.mapper.pojo.model.spi.PojoRawTypeModel;
 
-class JavaBeanTypeContextContainer implements JavaBeanSearchSessionTypeContextProvider {
+class JavaBeanTypeContextContainer implements JavaBeanSearchSessionTypeContextProvider, LoadingTypeContextProvider {
 
 	// Use a LinkedHashMap for deterministic iteration
 	private final Map<PojoRawTypeIdentifier<?>, JavaBeanIndexedTypeContext<?>> indexedTypeContexts = new LinkedHashMap<>();
@@ -49,6 +50,7 @@ public <E> AbstractJavaBeanTypeContext<E> forExactType(PojoRawTypeIdentifier<E>
 		return (AbstractJavaBeanTypeContext<E>) typeContexts.get( typeIdentifier );
 	}
 
+	@Override
 	@SuppressWarnings("unchecked")
 	public <E> JavaBeanIndexedTypeContext<E> indexedForExactClass(Class<E> clazz) {
 		return (JavaBeanIndexedTypeContext<E>) indexedTypeContextsByClass.get( clazz );

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/session/SearchSessionBuilder.java

@@ -6,8 +6,11 @@
  */
 package org.hibernate.search.mapper.javabean.session;
 
+import java.util.function.Consumer;
+
 import org.hibernate.search.engine.backend.work.execution.DocumentCommitStrategy;
 import org.hibernate.search.engine.backend.work.execution.DocumentRefreshStrategy;
+import org.hibernate.search.mapper.javabean.loading.LoadingOptions;
 
 public interface SearchSessionBuilder {
 
@@ -30,6 +33,12 @@ public interface SearchSessionBuilder {
 	 */
 	SearchSessionBuilder refreshStrategy(DocumentRefreshStrategy refreshStrategy);
 
+	/**
+	 * @param loadingOptionsContributor The default loading options.
+	 * @return {@code this} for method chaining.
+	 */
+	SearchSessionBuilder loading(Consumer<LoadingOptions> loadingOptionsContributor);
+
 	/**
 	 * @return The resulting session.
 	 */

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/session/impl/JavaBeanSearchSession.java

@@ -8,13 +8,15 @@
 
 import java.util.Collection;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
 
 import org.hibernate.search.engine.backend.work.execution.DocumentCommitStrategy;
 import org.hibernate.search.engine.backend.work.execution.DocumentRefreshStrategy;
 import org.hibernate.search.engine.backend.common.DocumentReference;
 import org.hibernate.search.engine.search.query.dsl.SearchQuerySelectStep;
 import org.hibernate.search.engine.backend.common.spi.DocumentReferenceConverter;
 import org.hibernate.search.mapper.javabean.common.EntityReference;
+import org.hibernate.search.mapper.javabean.loading.LoadingOptions;
 import org.hibernate.search.mapper.javabean.scope.SearchScope;
 import org.hibernate.search.mapper.javabean.scope.impl.SearchScopeImpl;
 import org.hibernate.search.mapper.javabean.loading.impl.JavaBeanSearchLoadingContext;
@@ -44,6 +46,8 @@ public class JavaBeanSearchSession extends AbstractPojoSearchSession<EntityRefer
 
 	private final DocumentCommitStrategy commitStrategy;
 	private final DocumentRefreshStrategy refreshStrategy;
+	private final Consumer<LoadingOptions> loadingOptionsContributor;
+
 	private SearchIndexingPlanImpl indexingPlan;
 	private SearchIndexer indexer;
 
@@ -54,6 +58,7 @@ private JavaBeanSearchSession(Builder builder) {
 		this.tenantId = builder.tenantId;
 		this.commitStrategy = builder.commitStrategy;
 		this.refreshStrategy = builder.refreshStrategy;
+		this.loadingOptionsContributor = builder.loadingOptionsContributor;
 	}
 
 	@Override
@@ -153,7 +158,11 @@ public PojoLoadingContext defaultLoadingContext() {
 	}
 
 	private PojoLoadingContextBuilder<?> loadingContextBuilder() {
-		return new JavaBeanSearchLoadingContext.Builder();
+		JavaBeanSearchLoadingContext.Builder builder = new JavaBeanSearchLoadingContext.Builder( typeContextProvider );
+		if ( loadingOptionsContributor != null ) {
+			loadingOptionsContributor.accept( builder );
+		}
+		return builder;
 	}
 
 	public static class Builder
@@ -163,6 +172,7 @@ public static class Builder
 		private String tenantId;
 		private DocumentCommitStrategy commitStrategy = DocumentCommitStrategy.FORCE;
 		private DocumentRefreshStrategy refreshStrategy = DocumentRefreshStrategy.NONE;
+		private Consumer<LoadingOptions> loadingOptionsContributor;
 
 		public Builder(JavaBeanSearchSessionMappingContext mappingContext,
 				JavaBeanSearchSessionTypeContextProvider typeContextProvider) {
@@ -188,6 +198,12 @@ public SearchSessionBuilder refreshStrategy(DocumentRefreshStrategy refreshStrat
 			return this;
 		}
 
+		@Override
+		public SearchSessionBuilder loading(Consumer<LoadingOptions> loadingOptionsContributor) {
+			this.loadingOptionsContributor = loadingOptionsContributor;
+			return this;
+		}
+
 		@Override
 		public JavaBeanSearchSession build() {
 			return new JavaBeanSearchSession( this );

mapper/javabean/src/main/java/org/hibernate/search/mapper/javabean/session/impl/JavaBeanSearchSessionTypeContextProvider.java

@@ -6,9 +6,11 @@
  */
 package org.hibernate.search.mapper.javabean.session.impl;
 
+import org.hibernate.search.mapper.javabean.loading.impl.LoadingTypeContextProvider;
 import org.hibernate.search.mapper.javabean.work.impl.SearchIndexingPlanTypeContextProvider;
 
-public interface JavaBeanSearchSessionTypeContextProvider extends SearchIndexingPlanTypeContextProvider {
+public interface JavaBeanSearchSessionTypeContextProvider
+		extends SearchIndexingPlanTypeContextProvider, LoadingTypeContextProvider {
 
 	JavaBeanSessionIndexedTypeContext<?> indexedForEntityName(String indexName);