A query on inherited beans can be restricted to the subtype of one of its children (#1602)

ENH: A query on inherited beans can be restricted to the subtype of one of its children

src/main/java/io/ebean/Query.java

@@ -1610,6 +1610,16 @@ public interface Query<T> {
    */
   Class<T> getBeanType();
 
+  /**
+   * Sets the inherit type. Must be a subtype of getBeanType
+   */
+  Query<T> setInheritType(Class<? extends T> type);
+
+  /**
+   * Returns the inherit type. This is normally the same as getBeanType() returns as long as no other type is set.
+   */
+  Class<? extends T> getInheritType();
+
   /**
    * Return the type of query being executed.
    */

src/main/java/io/ebeaninternal/server/querydefn/DefaultOrmQuery.java

@@ -47,6 +47,7 @@ import io.ebeaninternal.server.autotune.ProfilingListener;
 import io.ebeaninternal.server.core.SpiOrmQueryRequest;
 import io.ebeaninternal.server.deploy.BeanDescriptor;
 import io.ebeaninternal.server.deploy.BeanPropertyAssocMany;
+import io.ebeaninternal.server.deploy.InheritInfo;
 import io.ebeaninternal.server.deploy.TableJoin;
 import io.ebeaninternal.server.expression.DefaultExpressionList;
 import io.ebeaninternal.server.expression.SimpleExpression;
@@ -79,7 +80,9 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
 
   private final Class<T> beanType;
 
-  private final BeanDescriptor<T> beanDescriptor;
+  private final BeanDescriptor<T> rootBeanDescriptor;
+
+  private BeanDescriptor<T> beanDescriptor;
 
   private final SpiEbeanServer server;
 
@@ -276,6 +279,7 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
 
   public DefaultOrmQuery(BeanDescriptor<T> desc, SpiEbeanServer server, ExpressionFactory expressionFactory) {
     this.beanDescriptor = desc;
+    this.rootBeanDescriptor = desc;
     this.beanType = desc.getBeanType();
     this.server = server;
     this.orderById = server.getServerConfig().isDefaultOrderById();
@@ -1641,6 +1645,25 @@ public class DefaultOrmQuery<T> implements SpiQuery<T> {
     return beanType;
   }
 
+  @Override
+  public Class<? extends T> getInheritType() {
+    return beanDescriptor.getBeanType();
+  }
+
+  @Override
+  public Query<T> setInheritType(Class<? extends T> type) {
+    if (type == beanType) {
+      return this;
+    }
+    InheritInfo inheritInfo = rootBeanDescriptor.getInheritInfo();
+    inheritInfo = inheritInfo == null ? null : inheritInfo.readType(type);
+    if (inheritInfo == null) {
+      throw new IllegalArgumentException("Given type " + type + " is not a subtype of " + beanType);
+    }
+    beanDescriptor = (BeanDescriptor<T>) rootBeanDescriptor.getBeanDescriptor(type);
+    return this;
+  }
+
   @Override
   public String toString() {
     return "Query [" + whereExpressions + "]";

src/test/java/org/tests/inheritance/TestInheritanceQuery.java

@@ -0,0 +1,47 @@
+package org.tests.inheritance;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import org.tests.inherit.ChildA;
+import org.tests.inherit.ChildB;
+import org.tests.inherit.Parent;
+import org.junit.Test;
+
+import io.ebean.BaseTestCase;
+import io.ebean.Ebean;
+import io.ebean.Query;
+
+public class TestInheritanceQuery extends BaseTestCase {
+
+  @Test
+  public void test() {
+    Parent a = new ChildA(90, "Bean A");
+    Parent b1 = new ChildB(91, "Bean B1");
+    Parent b2 = new ChildB(92, "Bean B2");
+
+    server().save(a);
+    server().save(b1);
+    server().save(b2);
+
+    Query<Parent> query = Ebean.find(Parent.class);
+
+    query.where().in("val",90, 91); // restrict to a & b1
+
+    assertThat(query.findList()).hasSize(2); // a & b1
+
+    Query<Parent> query2 = query.copy();
+    query2.setInheritType(ChildA.class);
+
+    assertThat(query.findList()).hasSize(2);
+    assertThat(query2.findList()).containsExactly(a);
+
+    query2.setInheritType(ChildB.class);
+    assertThat(query2.findList()).containsExactly(b1);
+
+    query2 = Ebean.find(Parent.class).setInheritType(ChildB.class);
+    assertThat(query2.findList()).contains(b1, b2);
+
+    Ebean.delete(a);
+    Ebean.delete(b1);
+    Ebean.delete(b2);
+  }
+}