BatchFetchByJohnHuss.java

package jambalaya;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;

import org.apache.cayenne.Cayenne;
import org.apache.cayenne.DataObject;
import org.apache.cayenne.ObjectContext;
import org.apache.cayenne.PersistenceState;
import org.apache.cayenne.Persistent;
import org.apache.cayenne.access.ToManyList;
import org.apache.cayenne.exp.ExpressionFactory;
import org.apache.cayenne.exp.Property;
import org.apache.cayenne.map.ObjEntity;
import org.apache.cayenne.map.ObjRelationship;
import org.apache.cayenne.query.SelectQuery;

public class BatchFetchByJohnHuss {

	public static List<Persistent> batchFetch( Collection<? extends Persistent> sourceObjects, Property<?> path ) {
		List<List<Persistent>> result = batchFetch( sourceObjects, (List)Collections.singletonList( path ) );
		if( !result.isEmpty() ) {
			return result.get( 0 );
		}
		return Collections.emptyList();
	}

	/**
	* Like pre-fetching will resolve related objects in large batches rather
	than one at a time, but can be done after
	* the initial fetch has already occurred.
	*
	* The related objects must have only a single PK column. This restriction
	could be removed.
	*
	* This will issue a single query with a single IN expression for each
	path, so there may be limitations
	* on the number of sourceObjects it supports fetching at once, like 1024
	or less. This restriction could be removed.
	*
	* For to-many relationships there are restrictions: it doesn't work for
	nested paths or relationships without an inverse.
	*
	* @param sourceObjects
	* @param paths
	* @return all the objects that had to be fetched (this excludes objects
	that were not hollow to start with)
	*/
	public static List<List<Persistent>> batchFetch( Collection<? extends Persistent> sourceObjects, Collection<? extends Property<? extends Persistent>> paths ) {
		if( sourceObjects.isEmpty() ) {
			return Collections.emptyList();
		}
		sourceObjects = new ArrayList<>( new HashSet<Persistent>( sourceObjects ) ); // these have to be unique for the logic to work below
		ObjectContext context = sourceObjects.iterator().next().getObjectContext();
		List<List<Persistent>> result = new ArrayList<>( paths.size() );
		List<ObjEntity> entitiesForPaths = new ArrayList<>( paths.size() );
		ObjEntity objEntity = Cayenne.getObjEntity( sourceObjects.iterator().next() );

		for( Property<? extends Persistent> path : paths ) {
			ObjRelationship relationship = objEntity.getRelationship( path.getName() );
			if( relationship == null ) {
				continue;
			}
			if( relationship.isToMany() ) { // TODO: this doesn't work for paths or relationships without an inverse

				String reverseName = relationship.getReverseRelationshipName();
				if( reverseName == null ) {
					continue;
				}

				SelectQuery<Persistent> query = new SelectQuery<>( relationship.getTargetEntity(), ExpressionFactory.inExp( reverseName, sourceObjects ) );

				List<Persistent> matches = context.performQuery( query );
				result.add( matches );
				for( Persistent source : sourceObjects ) {
					DataObject sourceObject = (DataObject)source;
					ToManyList toManyList = new ToManyList( sourceObject, path.getName() );
					List<Persistent> relatedObjects = new ArrayList<>( ExpressionFactory.matchExp( reverseName, sourceObject ).filterObjects( matches ) );
					matches.removeAll( relatedObjects );
					toManyList.setValueDirectly( relatedObjects );
					sourceObject.writePropertyDirectly( path.getName(), toManyList );
				}
			}
			else {
				List<? extends Persistent> related = path.getFromAll( sourceObjects );
				related = new ArrayList<>( new HashSet<Persistent>( related ) ); // remove duplicates
				// remove objects that aren't hollow (probably none)
				Iterator<? extends Persistent> iterator = related.iterator();
				while( iterator.hasNext() ) {
					Persistent relatedObject = iterator.next();
					if( relatedObject == null || relatedObject.getPersistenceState() != PersistenceState.HOLLOW ) {
						iterator.remove();
					}
				}
				if( !related.isEmpty() ) {
					ObjEntity relatedEntity = Cayenne.getObjEntity( related.get( 0 ) );
					entitiesForPaths.add( relatedEntity );
					List<String> primaryKeyNames = new ArrayList<>( relatedEntity.getPrimaryKeyNames() );
					if( primaryKeyNames.size() != 1 ) {
						throw new IllegalArgumentException( "Cannot batch fetch. " + relatedEntity.getName() + " has multiple primary keys columns." );
					}
					String pk = primaryKeyNames.get( 0 );
					SelectQuery<Persistent> query = new SelectQuery<>( relatedEntity, ExpressionFactory.inDbExp( pk, related ) );
					result.add( context.performQuery( query ) );
				}
			}
		}
		return result;
	}
}