HHH-1775 - collection batch fetching

hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java

@@ -896,6 +896,9 @@ private Serializable getLoadedCollectionOwnerIdOrNull(CollectionEntry ce) {
 	public void addUninitializedCollection(CollectionPersister persister, PersistentCollection collection, Serializable id) {
 		CollectionEntry ce = new CollectionEntry(collection, persister, id, flushing);
 		addCollection(collection, ce, id);
+		if ( persister.getBatchSize() > 1 ) {
+			getBatchFetchQueue().addBatchLoadableCollection( collection, ce );
+		}
 	}
 
 	/**
@@ -905,6 +908,9 @@ public void addUninitializedCollection(CollectionPersister persister, Persistent
 	public void addUninitializedDetachedCollection(CollectionPersister persister, PersistentCollection collection) {
 		CollectionEntry ce = new CollectionEntry( persister, collection.getKey() );
 		addCollection( collection, ce, collection.getKey() );
+		if ( persister.getBatchSize() > 1 ) {
+			getBatchFetchQueue().addBatchLoadableCollection( collection, ce );
+		}
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/engine/spi/CollectionEntry.java

@@ -34,6 +34,7 @@
 import org.hibernate.AssertionFailure;
 import org.hibernate.HibernateException;
 import org.hibernate.MappingException;
+import org.hibernate.collection.internal.AbstractPersistentCollection;
 import org.hibernate.collection.spi.PersistentCollection;
 import org.hibernate.internal.CoreMessageLogger;
 import org.hibernate.persister.collection.CollectionPersister;
@@ -215,6 +216,12 @@ public void postInitialize(PersistentCollection collection) throws HibernateExce
 				collection.getSnapshot( getLoadedPersister() ) :
 				null;
 		collection.setSnapshot(loadedKey, role, snapshot);
+		if ( getLoadedPersister().getBatchSize() > 1 ) {
+			( (AbstractPersistentCollection) collection ).getSession()
+					.getPersistenceContext()
+					.getBatchFetchQueue()
+					.removeBatchLoadableCollection( this );
+		}
 	}
 
 	/**

hibernate-core/src/main/java/org/hibernate/event/internal/EvictVisitor.java

@@ -83,6 +83,9 @@ private void evictCollection(PersistentCollection collection) {
 							ce.getLoadedKey(),
 							getSession() ) );
 		}
+		if ( ce.getLoadedPersister() != null && ce.getLoadedPersister().getBatchSize() > 1 ) {
+			getSession().getPersistenceContext().getBatchFetchQueue().removeBatchLoadableCollection( ce );
+		}
 		if ( ce.getLoadedPersister() != null && ce.getLoadedKey() != null ) {
 			//TODO: is this 100% correct?
 			getSession().getPersistenceContext().getCollectionsByKey().remove(

hibernate-core/src/main/java/org/hibernate/persister/collection/AbstractCollectionPersister.java

@@ -1913,6 +1913,11 @@ public CollectionInitializer getInitializer() {
 		return initializer;
 	}
 
+	@Override
+	public int getBatchSize() {
+		return batchSize;
+	}
+
 	private class StandardOrderByAliasResolver implements OrderByAliasResolver {
 		private final String rootAlias;
 

hibernate-core/src/main/java/org/hibernate/persister/collection/CollectionPersister.java

@@ -307,4 +307,6 @@ public void insertRows(
 	public boolean indexExists(Serializable key, Object index, SessionImplementor session);
 	public boolean elementExists(Serializable key, Object element, SessionImplementor session);
 	public Object getElementByIndex(Serializable key, Object index, SessionImplementor session, Object owner);
+
+	public int getBatchSize();
 }

hibernate-core/src/main/java/org/hibernate/type/ManyToOneType.java

@@ -169,7 +169,7 @@ private void scheduleBatchLoadIfNeeded(Serializable id, SessionImplementor sessi
 		if ( uniqueKeyPropertyName == null && id != null ) {
 			final EntityPersister persister = session.getFactory().getEntityPersister( getAssociatedEntityName() );
 			final EntityKey entityKey = session.generateEntityKey( id, persister );
-			if ( !session.getPersistenceContext().containsEntity( entityKey ) ) {
+			if ( entityKey.isBatchLoadable() && !session.getPersistenceContext().containsEntity( entityKey ) ) {
 				session.getPersistenceContext().getBatchFetchQueue().addBatchLoadableEntityKey( entityKey );
 			}
 		}

hibernate-core/src/test/java/org/hibernate/test/cfg/persister/GoofyPersisterClassProvider.java

@@ -802,5 +802,10 @@ public boolean elementExists(Serializable key, Object element, SessionImplemento
 		public Object getElementByIndex(Serializable key, Object index, SessionImplementor session, Object owner) {
 			return null;  //To change body of implemented methods use File | Settings | File Templates.
 		}
+
+		@Override
+		public int getBatchSize() {
+			return 0;
+		}
 	}
 }