Simple helper to batch load of uninitialized LAZY relations in Hibernate entities.

BatchInitalizer.java

package org.xmedeko;

import java.util.List;
import java.util.Set;

import org.hibernate.annotations.BatchSize;
import org.hibernate.proxy.HibernateProxy;
import org.hibernate.proxy.LazyInitializer;

import com.google.common.base.Preconditions;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;

/**
 * Simple helper to batch load of uninitialized LAZY relations in Hibernate entities.
 * <p>
 * Check the Hibernate annotation {@link BatchSize} before using this class.
 * 
 * @param <K>
 *            Type of key.
 * 
 * @author xmedeko
 */
public abstract class BatchInitalizer<K> {

	private static final int BATCH_SIZE_DEFAULT = 1024;
	
	/** Max. number of items to load. Must be greater than 0. */
	private int batchSize = BATCH_SIZE_DEFAULT;

	/** Set of keys - IDs. */
	private final Set<K> keys;

	/**
	 * Constructor.
	 */
	public BatchLoader() {
		this(BATCH_SIZE_DEFAULT, 32);
	}

	/**
	 * Constructor.
	 * 
	 * @param batchSize
	 *            Max. number of items to load.
	 */
	public BatchLoader(int batchSize) {
		this(batchSize, batchSize);
	}

	/**
	 * Constructor.
	 * 
	 * @param batchSize
	 *            Max. number of items to load. If less or equal to zero, then set the default.
	 * @param expectedSize
	 *            Expected kesy size. Just for the optimalization. If not set, set the same as batchSize.
	 */
	public BatchLoader(int batchSize, int expectedSize) {
		if (batchSize <= 0) {
			batchSize = BATCH_SIZE_DEFAULT;
		}
		setBatchSize(batchSize);
		
		if (expectedSize <= 0) {
			expectedSize = batchSize;
		}
		keys = Sets.newHashSetWithExpectedSize(expectedSize);
	}

	public int getBatchSize() {
		return batchSize;
	}

	public BatchLoader<K> setBatchSize(int batchSize) {
		Preconditions.checkArgument(batchSize > 0);
		this.batchSize = batchSize;
		return this;
	}

	/**
	 * Add a key.
	 * 
	 * @param key
	 */
	public void addKey(K key) {
		keys.add(key);
	}

	/**
	 * Add a key - ID, when the Hibernate entity has not been initialized yet.
	 * 
	 * @param entity
	 */
	public void addUninitializedKey(Object entity) {
		if (!(entity instanceof HibernateProxy)) {
			return;
		}
		LazyInitializer lazyInitializer = ((HibernateProxy) entity).getHibernateLazyInitializer();
		if (!lazyInitializer.isUninitialized()) {
			return;
		}
		@SuppressWarnings("unchecked")
		K id = (K) lazyInitializer.getIdentifier();
		addKey(id);
	}

	/**
	 * Load the entities by keys.
	 */
	public void load() {
		Iterable<List<K>> partitions = Iterables.partition(keys, batchSize);
		for (List<K> partition : partitions) {
			loadAll(partition);
		}

		// remove the loaded keys
		keys.clear();
	}

	/**
	 * Batch load the entities by keys.
	 * 
	 * @param keys
	 *            keys.
	 */
	public abstract void loadAll(List<K> keys);

}