OGM-767 Copy batching classes from ORM as is

core/src/main/java/org/hibernate/ogm/loader/entity/impl/BatchingEntityLoader.java

@@ -0,0 +1,112 @@
+/*
+ * Hibernate OGM, Domain model persistence for NoSQL datastores
+ *
+ * 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.ogm.loader.entity.impl;
+
+import java.io.Serializable;
+import java.sql.SQLException;
+import java.util.Arrays;
+import java.util.List;
+
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.QueryParameters;
+import org.hibernate.engine.spi.SessionImplementor;
+import org.hibernate.loader.Loader;
+import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.pretty.MessageHelper;
+import org.hibernate.type.Type;
+
+import org.jboss.logging.Logger;
+
+/**
+ * The base contract for loaders capable of performing batch-fetch loading of entities using multiple primary key
+ * values in the SQL <tt>WHERE</tt> clause.
+ *
+ * @author Gavin King
+ * @author Steve Ebersole
+ *
+ * @see BatchingEntityLoaderBuilder
+ * @see UniqueEntityLoader
+ */
+public abstract class BatchingEntityLoader implements UniqueEntityLoader {
+	private static final Logger log = Logger.getLogger( BatchingEntityLoader.class );
+
+	private final EntityPersister persister;
+
+	public BatchingEntityLoader(EntityPersister persister) {
+		this.persister = persister;
+	}
+
+	public EntityPersister persister() {
+		return persister;
+	}
+
+	@Override
+	@Deprecated
+	public Object load(Serializable id, Object optionalObject, SessionImplementor session) {
+		return load( id, optionalObject, session, LockOptions.NONE );
+	}
+
+	protected QueryParameters buildQueryParameters(
+			Serializable id,
+			Serializable[] ids,
+			Object optionalObject,
+			LockOptions lockOptions) {
+		Type[] types = new Type[ids.length];
+		Arrays.fill( types, persister().getIdentifierType() );
+
+		QueryParameters qp = new QueryParameters();
+		qp.setPositionalParameterTypes( types );
+		qp.setPositionalParameterValues( ids );
+		qp.setOptionalObject( optionalObject );
+		qp.setOptionalEntityName( persister().getEntityName() );
+		qp.setOptionalId( id );
+		qp.setLockOptions( lockOptions );
+		return qp;
+	}
+
+	protected Object getObjectFromList(List results, Serializable id, SessionImplementor session) {
+		for ( Object obj : results ) {
+			final boolean equal = persister.getIdentifierType().isEqual(
+					id,
+					session.getContextEntityIdentifier( obj ),
+					session.getFactory()
+			);
+			if ( equal ) {
+				return obj;
+			}
+		}
+		return null;
+	}
+
+	protected Object doBatchLoad(
+			Serializable id,
+			Loader loaderToUse,
+			SessionImplementor session,
+			Serializable[] ids,
+			Object optionalObject,
+			LockOptions lockOptions) {
+		if ( log.isDebugEnabled() ) {
+			log.debugf( "Batch loading entity: %s", MessageHelper.infoString( persister, ids, session.getFactory() ) );
+		}
+
+		QueryParameters qp = buildQueryParameters( id, ids, optionalObject, lockOptions );
+
+		try {
+			final List results = loaderToUse.doQueryAndInitializeNonLazyCollections( session, qp, false );
+			log.debug( "Done entity batch load" );
+			return getObjectFromList(results, id, session);
+		}
+		catch ( SQLException sqle ) {
+			throw session.getFactory().getSQLExceptionHelper().convert(
+					sqle,
+					"could not load an entity batch: " + MessageHelper.infoString( persister(), ids, session.getFactory() ),
+					loaderToUse.getSQLString()
+			);
+		}
+	}
+
+}

core/src/main/java/org/hibernate/ogm/loader/entity/impl/BatchingEntityLoaderBuilder.java

@@ -0,0 +1,117 @@
+/*
+ * Hibernate OGM, Domain model persistence for NoSQL datastores
+ *
+ * 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.ogm.loader.entity.impl;
+
+import org.hibernate.LockMode;
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.LoadQueryInfluencers;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.persister.entity.OuterJoinLoadable;
+
+/**
+ * The contract for building {@link UniqueEntityLoader} capable of performing batch-fetch loading.  Intention
+ * is to build these instances, by first calling the static {@link #getBuilder}, and then calling the appropriate
+ * {@link #buildLoader} method.
+ *
+ * @author Steve Ebersole
+ *
+ * @see org.hibernate.loader.BatchFetchStyle
+ */
+public abstract class BatchingEntityLoaderBuilder {
+	public static BatchingEntityLoaderBuilder getBuilder(SessionFactoryImplementor factory) {
+		switch ( factory.getSettings().getBatchFetchStyle() ) {
+			case PADDED: {
+				return PaddedBatchingEntityLoaderBuilder.INSTANCE;
+			}
+			case DYNAMIC: {
+				return DynamicBatchingEntityLoaderBuilder.INSTANCE;
+			}
+			default: {
+				return org.hibernate.loader.entity.plan.LegacyBatchingEntityLoaderBuilder.INSTANCE;
+//				return LegacyBatchingEntityLoaderBuilder.INSTANCE;
+			}
+		}
+	}
+
+	/**
+	 * Builds a batch-fetch capable loader based on the given persister, lock-mode, etc.
+	 *
+	 * @param persister The entity persister
+	 * @param batchSize The maximum number of ids to batch-fetch at once
+	 * @param lockMode The lock mode
+	 * @param factory The SessionFactory
+	 * @param influencers Any influencers that should affect the built query
+	 *
+	 * @return The loader.
+	 */
+	public UniqueEntityLoader buildLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockMode lockMode,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		if ( batchSize <= 1 ) {
+			// no batching
+			return buildNonBatchingLoader( persister, lockMode, factory, influencers );
+		}
+		return buildBatchingLoader( persister, batchSize, lockMode, factory, influencers );
+	}
+
+	protected UniqueEntityLoader buildNonBatchingLoader(
+			OuterJoinLoadable persister,
+			LockMode lockMode,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		return new EntityLoader( persister, lockMode, factory, influencers );
+	}
+
+	protected abstract UniqueEntityLoader buildBatchingLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockMode lockMode,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers);
+
+	/**
+	 * Builds a batch-fetch capable loader based on the given persister, lock-options, etc.
+	 *
+	 * @param persister The entity persister
+	 * @param batchSize The maximum number of ids to batch-fetch at once
+	 * @param lockOptions The lock options
+	 * @param factory The SessionFactory
+	 * @param influencers Any influencers that should affect the built query
+	 *
+	 * @return The loader.
+	 */
+	public UniqueEntityLoader buildLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockOptions lockOptions,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		if ( batchSize <= 1 ) {
+			// no batching
+			return buildNonBatchingLoader( persister, lockOptions, factory, influencers );
+		}
+		return buildBatchingLoader( persister, batchSize, lockOptions, factory, influencers );
+	}
+
+	protected UniqueEntityLoader buildNonBatchingLoader(
+			OuterJoinLoadable persister,
+			LockOptions lockOptions,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		return new EntityLoader( persister, lockOptions, factory, influencers );
+	}
+
+	protected abstract UniqueEntityLoader buildBatchingLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockOptions lockOptions,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers);
+}

core/src/main/java/org/hibernate/ogm/loader/entity/impl/PaddedBatchingEntityLoaderBuilder.java

@@ -0,0 +1,123 @@
+/*
+ * Hibernate OGM, Domain model persistence for NoSQL datastores
+ *
+ * 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.ogm.loader.entity.impl;
+
+import java.io.Serializable;
+
+import org.hibernate.HibernateException;
+import org.hibernate.LockMode;
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.LoadQueryInfluencers;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.engine.spi.SessionImplementor;
+import org.hibernate.internal.util.collections.ArrayHelper;
+import org.hibernate.loader.Loader;
+import org.hibernate.persister.entity.OuterJoinLoadable;
+
+/**
+ * @author Steve Ebersole
+ */
+class PaddedBatchingEntityLoaderBuilder extends BatchingEntityLoaderBuilder {
+	public static final PaddedBatchingEntityLoaderBuilder INSTANCE = new PaddedBatchingEntityLoaderBuilder();
+
+	@Override
+	protected UniqueEntityLoader buildBatchingLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockMode lockMode,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		return new PaddedBatchingEntityLoader( persister, batchSize, lockMode, factory, influencers );
+	}
+
+	@Override
+	protected UniqueEntityLoader buildBatchingLoader(
+			OuterJoinLoadable persister,
+			int batchSize,
+			LockOptions lockOptions,
+			SessionFactoryImplementor factory,
+			LoadQueryInfluencers influencers) {
+		return new PaddedBatchingEntityLoader( persister, batchSize, lockOptions, factory, influencers );
+	}
+
+	public static class PaddedBatchingEntityLoader extends BatchingEntityLoader {
+		private final int[] batchSizes;
+		private final Loader[] loaders;
+
+		public PaddedBatchingEntityLoader(
+				OuterJoinLoadable persister,
+				int maxBatchSize,
+				LockMode lockMode,
+				SessionFactoryImplementor factory,
+				LoadQueryInfluencers loadQueryInfluencers) {
+			super( persister );
+			this.batchSizes = ArrayHelper.getBatchSizes( maxBatchSize );
+			this.loaders = new Loader[ batchSizes.length ];
+			for ( int i = 0; i < batchSizes.length; i++ ) {
+				this.loaders[i] = new EntityLoader( persister, batchSizes[i], lockMode, factory, loadQueryInfluencers);
+			}
+			validate( maxBatchSize );
+		}
+
+		private void validate(int max) {
+			// these are more indicative of internal problems then user error...
+			if ( batchSizes[0] != max ) {
+				throw new HibernateException( "Unexpected batch size spread" );
+			}
+			if ( batchSizes[batchSizes.length-1] != 1 ) {
+				throw new HibernateException( "Unexpected batch size spread" );
+			}
+		}
+
+		public PaddedBatchingEntityLoader(
+				OuterJoinLoadable persister,
+				int maxBatchSize,
+				LockOptions lockOptions,
+				SessionFactoryImplementor factory,
+				LoadQueryInfluencers loadQueryInfluencers) {
+			super( persister );
+			this.batchSizes = ArrayHelper.getBatchSizes( maxBatchSize );
+			this.loaders = new Loader[ batchSizes.length ];
+			for ( int i = 0; i < batchSizes.length; i++ ) {
+				this.loaders[i] = new EntityLoader( persister, batchSizes[i], lockOptions, factory, loadQueryInfluencers);
+			}
+			validate( maxBatchSize );
+		}
+
+		@Override
+		public Object load(Serializable id, Object optionalObject, SessionImplementor session, LockOptions lockOptions) {
+			final Serializable[] batch = session.getPersistenceContext()
+					.getBatchFetchQueue()
+					.getEntityBatch( persister(), id, batchSizes[0], persister().getEntityMode() );
+
+			final int numberOfIds = ArrayHelper.countNonNull( batch );
+			if ( numberOfIds <= 1 ) {
+				return ( (UniqueEntityLoader) loaders[batchSizes.length-1] ).load( id, optionalObject, session );
+			}
+
+			// Uses the first batch-size bigger than the number of actual ids in the batch
+			int indexToUse = batchSizes.length-1;
+			for ( int i = 0; i < batchSizes.length-1; i++ ) {
+				if ( batchSizes[i] >= numberOfIds ) {
+					indexToUse = i;
+				}
+				else {
+					break;
+				}
+			}
+
+			final Serializable[] idsToLoad = new Serializable[ batchSizes[indexToUse] ];
+			System.arraycopy( batch, 0, idsToLoad, 0, numberOfIds );
+			for ( int i = numberOfIds; i < batchSizes[indexToUse]; i++ ) {
+				idsToLoad[i] = id;
+			}
+
+			return doBatchLoad( id, loaders[indexToUse], session, idsToLoad, optionalObject, lockOptions );
+		}
+	}
+
+}