From 0efbe061ab88fea91ac310de134dff6df5e75563 Mon Sep 17 00:00:00 2001
From: Robin Bygrave <robin.bygrave@gmail.com>
Date: Tue, 24 Nov 2015 01:01:37 +1300
Subject: [PATCH] #461 - ENH: Add @Draftable support - optimise publish query

---
 .../server/deploy/BeanDescriptor.java         |  8 +++++++
 .../deploy/BeanDescriptorDraftHelp.java       | 22 +++++++++++++++++++
 .../server/persist/DefaultPersister.java      |  8 ++++---
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptor.java b/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptor.java
index 657989df1..5076cc767 100644
--- a/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptor.java
+++ b/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptor.java
@@ -1,6 +1,7 @@
 package com.avaje.ebeaninternal.server.deploy;
 
 import com.avaje.ebean.OrderBy;
+import com.avaje.ebean.Query;
 import com.avaje.ebean.SqlUpdate;
 import com.avaje.ebean.Transaction;
 import com.avaje.ebean.ValuePair;
@@ -1941,6 +1942,13 @@ public class BeanDescriptor<T> implements MetaBeanInfo, SpiBeanType<T> {
     }
   }
 
+  /**
+   * Optimise the draft query fetching any draftable element relationships.
+   */
+  public void draftQueryOptimise(Query<T> query) {
+    draftHelp.draftQueryOptimise(query);
+  }
+
   /**
    * Return true if this entity bean has history support.
    */
diff --git a/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptorDraftHelp.java b/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptorDraftHelp.java
index 6f62f6185..118c8fdf4 100644
--- a/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptorDraftHelp.java
+++ b/src/main/java/com/avaje/ebeaninternal/server/deploy/BeanDescriptorDraftHelp.java
@@ -1,5 +1,6 @@
 package com.avaje.ebeaninternal.server.deploy;
 
+import com.avaje.ebean.Query;
 import com.avaje.ebean.bean.EntityBean;
 
 /**
@@ -49,4 +50,25 @@ public final class BeanDescriptorDraftHelp<T> {
 
     return liveBean;
   }
+
+  /**
+   * Fetch draftable element relationships.
+   */
+  public void draftQueryOptimise(Query<T> query) {
+
+    BeanPropertyAssocOne<?>[] one = desc.propertiesOne();
+    for (int i = 0; i < one.length; i++) {
+      if (one[i].getTargetDescriptor().isDraftableElement()) {
+        query.fetch(one[i].getName());
+      }
+    }
+
+    BeanPropertyAssocMany<?>[] many = desc.propertiesMany();
+    for (int i = 0; i < many.length; i++) {
+      if (many[i].getTargetDescriptor().isDraftableElement()) {
+        query.fetch(many[i].getName());
+      }
+    }
+
+  }
 }
diff --git a/src/main/java/com/avaje/ebeaninternal/server/persist/DefaultPersister.java b/src/main/java/com/avaje/ebeaninternal/server/persist/DefaultPersister.java
index 8b110a6ff..884395cac 100644
--- a/src/main/java/com/avaje/ebeaninternal/server/persist/DefaultPersister.java
+++ b/src/main/java/com/avaje/ebeaninternal/server/persist/DefaultPersister.java
@@ -134,6 +134,7 @@ public final class DefaultPersister implements Persister {
 
     Class<T> beanType = query.getBeanType();
     BeanDescriptor<T> desc = server.getBeanDescriptor(beanType);
+    desc.draftQueryOptimise(query);
 
     List<T> draftBeans = server.findList(query, transaction);
     PUB.debug("publish [{}] count[{}]", desc.getName(), draftBeans.size());
@@ -149,9 +150,10 @@ public final class DefaultPersister implements Persister {
     }
 
     // fetch existing live beans to update (or insert if missing)
-    Map<?, T> liveBeans = server.find(beanType)
-        .where().idIn(idList)
-        .findMap();
+    Query<T> liveBeansQuery = server.find(beanType).where().idIn(idList).query();
+    desc.draftQueryOptimise(liveBeansQuery);
+
+    Map<?, T> liveBeans = liveBeansQuery.findMap();
 
     List<T> livePublish = new ArrayList<T>(idList.size());