Fix for #228 - LazyLoad bug with QueryIterator. Getting EntityNotFoundException("Bean has been deleted - lazy loading failed")

src/main/java/com/avaje/ebean/bean/BeanCollection.java

@@ -83,9 +83,17 @@ public interface BeanCollection<E> extends Serializable {
   public void setBeanCollectionTouched(BeanCollectionTouched notify);
 
   /**
-   * Set the loader that will be used to lazy/query load this collection.
+   * Return true if the collection has been registered with the batch loading context.
    */
-  public void setLoader(int beanLoaderIndex, BeanCollectionLoader beanLoader);
+  public boolean isRegisteredWithLoadContext();
+
+  /**
+   * Set the loader that will be used to lazy/query load this collection.
+   * <p>
+   * This is effectively the batch loading context this collection is registered with.
+   * </p>
+   */
+  public void setLoader(BeanCollectionLoader beanLoader);
 
   /**
    * Set to true if you want the BeanCollection to be treated as read only. This

src/main/java/com/avaje/ebean/common/AbstractBeanCollection.java

@@ -29,7 +29,10 @@ public abstract class AbstractBeanCollection<E> implements BeanCollection<E> {
 
   protected transient ExpressionList<?> filterMany;
 
-  protected int loaderIndex;
+  /**
+   * Flag set when registered with the batch loading context.
+   */
+  protected boolean registeredWithLoadContext;
 
   protected String ebeanServerName;
 
@@ -125,8 +128,12 @@ public abstract class AbstractBeanCollection<E> implements BeanCollection<E> {
     this.beanCollectionTouched = notify;
   }
 
-  public void setLoader(int beanLoaderIndex, BeanCollectionLoader loader) {
-    this.loaderIndex = beanLoaderIndex;
+  public boolean isRegisteredWithLoadContext() {
+    return registeredWithLoadContext;
+  }
+
+  public void setLoader(BeanCollectionLoader loader) {
+    this.registeredWithLoadContext = true;
     this.loader = loader;
     this.ebeanServerName = loader.getName();
   }

src/main/java/com/avaje/ebeaninternal/api/LoadContext.java

@@ -44,8 +44,11 @@ public interface LoadContext {
 
 	/**
 	 * Set the persistence context used by this query and future lazy loading.
+	 * <p>
+	 * Used by query iterator when processing large result sets.
+	 * </p>
 	 */
-	public void setPersistenceContext(PersistenceContext persistenceContext);
+	public void resetPersistenceContext(PersistenceContext persistenceContext);
 
 	/**
 	 * Register a Bean for lazy loading.

src/main/java/com/avaje/ebeaninternal/server/core/OrmQueryRequest.java

@@ -181,7 +181,7 @@ public final class OrmQueryRequest<T> extends BeanRequest implements BeanQueryRe
    */
   public void flushPersistenceContextOnIterate() {
     persistenceContext = new DefaultPersistenceContext();
-    loadContext.setPersistenceContext(persistenceContext);
+    loadContext.resetPersistenceContext(persistenceContext);
   }
 
   /**

src/main/java/com/avaje/ebeaninternal/server/loadcontext/DLoadBaseContext.java

@@ -100,6 +100,6 @@ public abstract class DLoadBaseContext {
   
   protected PersistenceContext getPersistenceContext() {
     return parent.getPersistenceContext();
-  }  
+  }
 
 }

src/main/java/com/avaje/ebeaninternal/server/loadcontext/DLoadBeanContext.java

@@ -17,7 +17,6 @@ import com.avaje.ebeaninternal.server.querydefn.OrmQueryProperties;
 
 /**
  * Default implementation of LoadBeanContext.
- *
  */
 public class DLoadBeanContext extends DLoadBaseContext implements LoadBeanContext{
 	
@@ -33,7 +32,17 @@ public class DLoadBeanContext extends DLoadBaseContext implements LoadBeanContex
     this.bufferList = (!queryFetch) ? null : new ArrayList<DLoadBeanContext.LoadBuffer>();
     this.currentBuffer = createBuffer(firstBatchSize);  
   }
-  
+
+  /**
+   * Reset the buffers after a query iterator reset.
+   */
+  public void clear() {
+    if (bufferList != null) {
+      bufferList.clear();
+    }
+    currentBuffer = createBuffer(secondaryBatchSize);
+  }
+
   protected void configureQuery(SpiQuery<?> query, String lazyLoadProperty) {
 
     // propagate the readOnly state

src/main/java/com/avaje/ebeaninternal/server/loadcontext/DLoadContext.java

@@ -1,14 +1,6 @@
 package com.avaje.ebeaninternal.server.loadcontext;
 
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-
-import com.avaje.ebean.bean.BeanCollection;
-import com.avaje.ebean.bean.EntityBeanIntercept;
-import com.avaje.ebean.bean.ObjectGraphNode;
-import com.avaje.ebean.bean.ObjectGraphOrigin;
-import com.avaje.ebean.bean.PersistenceContext;
+import com.avaje.ebean.bean.*;
 import com.avaje.ebeaninternal.api.LoadContext;
 import com.avaje.ebeaninternal.api.LoadSecondaryQuery;
 import com.avaje.ebeaninternal.api.SpiEbeanServer;
@@ -21,6 +13,10 @@ import com.avaje.ebeaninternal.server.deploy.BeanPropertyAssocMany;
 import com.avaje.ebeaninternal.server.el.ElPropertyValue;
 import com.avaje.ebeaninternal.server.querydefn.OrmQueryProperties;
 
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 /**
  * Default implementation of LoadContext.
  */
@@ -220,8 +216,16 @@ public class DLoadContext implements LoadContext {
 		return persistenceContext;
 	}
 
-	public void setPersistenceContext(PersistenceContext persistenceContext) {
+	public void resetPersistenceContext(PersistenceContext persistenceContext) {
 		this.persistenceContext = persistenceContext;
+		// clear the load contexts for beans and beanCollections
+		for (DLoadBeanContext beanContext : beanMap.values()) {
+			beanContext.clear();
+		}
+		for (DLoadManyContext manyContext : manyMap.values()) {
+			manyContext.clear();
+		}
+		this.rootBeanContext.clear();
 	}
 
 	public void register(String path, EntityBeanIntercept ebi){

src/main/java/com/avaje/ebeaninternal/server/loadcontext/DLoadManyContext.java

@@ -43,7 +43,17 @@ public class DLoadManyContext extends DLoadBaseContext implements LoadManyContex
     }
     return buffer;
   }
- 
+
+  /**
+   * Reset the buffers for a query iterator reset.
+   */
+  public void clear() {
+    if (bufferList != null) {
+      bufferList.clear();
+    }
+    currentBuffer = createBuffer(secondaryBatchSize);
+  }
+
   public void configureQuery(SpiQuery<?> query){
 		
 		// propagate the readOnly state
@@ -80,7 +90,7 @@ public class DLoadManyContext extends DLoadBaseContext implements LoadManyContex
       currentBuffer = createBuffer(secondaryBatchSize);
 		}
 		currentBuffer.add(bc);
-    bc.setLoader(0, currentBuffer);
+    bc.setLoader(currentBuffer);
 	}
 
   public void loadSecondaryQuery(OrmQueryRequest<?> parentRequest) {

src/main/java/com/avaje/ebeaninternal/server/query/SqlTreeNodeBean.java

@@ -345,7 +345,7 @@ public class SqlTreeNodeBean implements SqlTreeNode {
       } else {
         // create a proxy for the many (deferred fetching)
         BeanCollection<?> ref = manys[i].createReferenceIfNull(localBean);
-        if (ref != null) {
+        if (ref != null && !ref.isRegisteredWithLoadContext()) {
           ctx.register(manys[i].getName(), ref);
         }
       }

src/test/java/com/avaje/tests/query/TestQueryFindIterate.java

@@ -2,17 +2,20 @@ package com.avaje.tests.query;
 
 import javax.persistence.PersistenceException;
 
-import org.junit.Assert;
+import com.avaje.ebean.*;
+import com.avaje.tests.model.basic.Order;
+import com.avaje.tests.model.basic.OrderShipment;
+import org.avaje.ebeantest.LoggedSqlCollector;
 import org.junit.Test;
 
-import com.avaje.ebean.BaseTestCase;
-import com.avaje.ebean.Ebean;
-import com.avaje.ebean.EbeanServer;
-import com.avaje.ebean.Query;
-import com.avaje.ebean.QueryIterator;
 import com.avaje.tests.model.basic.Customer;
 import com.avaje.tests.model.basic.ResetBasicData;
 
+import java.util.List;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
 public class TestQueryFindIterate extends BaseTestCase {
 
   @Test
@@ -40,9 +43,108 @@ public class TestQueryFindIterate extends BaseTestCase {
       it.close();
     }
 
-    Assert.assertEquals(2, count);
+    assertEquals(2, count);
   }
-  
+
+  @Test
+  public void testWithLazyLoading() {
+
+    ResetBasicData.reset();
+
+    QueryIterator<Order> queryIterator = Ebean.find(Order.class)
+            //.select("orderDate")
+            .where().gt("id",0).le("id",10)
+            .findIterate();
+
+    try {
+      while (queryIterator.hasNext()) {
+        Order order = queryIterator.next();
+        Customer customer = order.getCustomer();
+        // invoke lazy loading on customer, order details and order shipments
+        System.out.println("order: " + order.getId() + " customerName:" + customer.getName()+" details:"+order.getDetails().size()+" shipments:"+order.getShipments().size());
+      }
+
+    } finally {
+      queryIterator.close();
+    }
+
+  }
+
+  @Test
+  public void testWithLazyBatchSize() {
+
+    ResetBasicData.reset();
+
+    LoggedSqlCollector.start();
+
+    QueryIterator<Order> queryIterator = Ebean.find(Order.class)
+            .setLazyLoadBatchSize(10)
+            .select("status, orderDate")
+            .fetch("customer", "name")
+            .where().gt("id",0).le("id",10)
+            .findIterate();
+
+    try {
+      while (queryIterator.hasNext()) {
+        Order order = queryIterator.next();
+        Customer customer = order.getCustomer();
+        System.out.println("order: " + order.getId() + " customerName:" + customer.getName()+" details:"+order.getDetails().size()+" shipments:"+order.getShipments().size());
+      }
+
+    } finally {
+      queryIterator.close();
+    }
+
+    List<String> loggedSql = LoggedSqlCollector.stop();
+
+    assertEquals(3, loggedSql.size());
+    assertTrue(loggedSql.get(0).contains("select t0.id c0, t0.status c1, t0.order_date c2, t1.id c3, t1.name c4 from o_order t0 join o_customer t1"));
+    assertTrue(loggedSql.get(1).contains("select t0.order_id c0, t0.id c1, t0.order_qty c2, t0.ship_qty c3, t0.unit_price c4"));
+    assertTrue(loggedSql.get(2).contains("select t0.order_id c0, t0.id c1, t0.ship_time c2, t0.cretime c3, t0.updtime c4, t0.order_id c5 from or_order_ship"));
+  }
+
+  @Test
+  public void testWithTwoJoins() {
+
+    ResetBasicData.reset();
+
+    LoggedSqlCollector.start();
+
+    // make sure we don't hit the L2 cache for order shipments
+    Ebean.getServerCacheManager().getBeanCache(Order.class).clear();
+    Ebean.getServerCacheManager().getQueryCache(Order.class).clear();
+    Ebean.getServerCacheManager().getBeanCache(OrderShipment.class).clear();
+    Ebean.getServerCacheManager().getQueryCache(OrderShipment.class).clear();
+
+    QueryIterator<Order> queryIterator = Ebean.find(Order.class)
+            .setLazyLoadBatchSize(10)
+            .setUseCache(false)
+            .setUseQueryCache(false)
+            .select("status, orderDate")
+            .fetch("customer", "name")
+            .fetch("details")
+            .where().gt("id",0).le("id",10)
+            .order().asc("id")
+            .findIterate();
+
+    try {
+      while (queryIterator.hasNext()) {
+        Order order = queryIterator.next();
+        Customer customer = order.getCustomer();
+        System.out.println("order: " + order.getId() + " customerName:" + customer.getName()+" details:"+order.getDetails().size()+" shipments:"+order.getShipments().size());
+      }
+
+    } finally {
+      queryIterator.close();
+    }
+
+    List<String> loggedSql = LoggedSqlCollector.stop();
+
+    assertEquals("Got SQL: "+loggedSql, 2, loggedSql.size());
+    assertTrue(loggedSql.get(0).contains("select t0.id c0, t0.status c1, t0.order_date c2, t1.id c3, t1.name c4, t2.id c5, t2.order_qty c6, t2.ship_qty"));
+    assertTrue(loggedSql.get(1).contains("select t0.order_id c0, t0.id c1, t0.ship_time c2, t0.cretime c3, t0.updtime c4, t0.order_id c5 from or_order_ship"));
+  }
+
   @Test(expected=PersistenceException.class)
   public void testWithExceptionInQuery() {
 
@@ -59,7 +161,7 @@ public class TestQueryFindIterate extends BaseTestCase {
     // this throws an exception immediately
     QueryIterator<Customer> it = query.findIterate();
     it.hashCode();
-    Assert.assertTrue("Never get here as exception thrown", false);
+    assertTrue("Never get here as exception thrown", false);
   }