HHH-9411 : Fix downcast of treat statement.

hibernate-core/src/main/java/org/hibernate/persister/entity/SingleTableEntityPersister.java

@@ -620,6 +620,21 @@ public class SingleTableEntityPersister extends AbstractEntityPersister {
 			if ( !queryable.isAbstract() ) {
 				values.add( queryable.getDiscriminatorSQLValue() );
 			}
+			else if ( queryable.hasSubclasses() ) {
+				// if the treat is an abstract class, add the concrete implementations to values if any
+				Set<String> actualSubClasses = queryable.getEntityMetamodel().getSubclassEntityNames();
+
+				for ( String actualSubClass : actualSubClasses ) {
+					if ( actualSubClass.equals( subclass ) ) {
+						continue;
+					}
+
+					Queryable actualQueryable = (Queryable) getFactory().getEntityPersister( actualSubClass );
+					if ( !actualQueryable.hasSubclasses() ) {
+						values.add( actualQueryable.getDiscriminatorSQLValue() );
+					}
+				}
+			}
 		}
 		return values.toArray( new String[values.size()] );
 	}

hibernate-core/src/test/java/org/hibernate/test/jpa/ql/TreatKeywordTest.java

@@ -11,14 +11,17 @@ import javax.persistence.DiscriminatorType;
 import javax.persistence.DiscriminatorValue;
 import javax.persistence.Entity;
 import javax.persistence.FetchType;
+import javax.persistence.GeneratedValue;
 import javax.persistence.Id;
 import javax.persistence.Inheritance;
 import javax.persistence.InheritanceType;
 import javax.persistence.ManyToOne;
 import javax.persistence.Table;
+import java.util.Arrays;
 import java.util.List;
 
 import org.hibernate.Session;
+import org.hibernate.Transaction;
 
 import org.junit.Test;
 
@@ -38,7 +41,8 @@ public class TreatKeywordTest extends BaseCoreFunctionalTestCase {
 		return new Class[] {
 				JoinedEntity.class, JoinedEntitySubclass.class, JoinedEntitySubSubclass.class,
 				JoinedEntitySubclass2.class, JoinedEntitySubSubclass2.class,
-				DiscriminatorEntity.class, DiscriminatorEntitySubclass.class, DiscriminatorEntitySubSubclass.class
+				DiscriminatorEntity.class, DiscriminatorEntitySubclass.class, DiscriminatorEntitySubSubclass.class,
+				Animal.class, Dog.class, Dachshund.class, Greyhound.class
 		};
 	}
 
@@ -184,6 +188,25 @@ public class TreatKeywordTest extends BaseCoreFunctionalTestCase {
 		s.close();
 	}
 
+	@Test
+	@TestForIssue(jiraKey = "HHH-9411")
+	public void testTreatWithRestrictionOnAbstractClass() {
+		Session s = openSession();
+		Transaction tx = s.beginTransaction();
+
+		Greyhound greyhound = new Greyhound();
+		Dachshund dachshund = new Dachshund();
+		s.save( greyhound );
+		s.save( dachshund );
+
+		List results = s.createQuery( "select treat (a as Dog) from Animal a where a.fast = TRUE" ).list();
+
+		assertEquals( Arrays.asList( greyhound ), results );
+
+		tx.commit();
+		s.close();
+	}
+
 	@Entity( name = "JoinedEntity" )
 	@Table( name = "JoinedEntity" )
 	@Inheritance( strategy = InheritanceType.JOINED )
@@ -334,4 +357,38 @@ public class TreatKeywordTest extends BaseCoreFunctionalTestCase {
 			super( id, name, other );
 		}
 	}
+
+	@Entity(name = "Animal")
+	public static abstract class Animal {
+		@Id
+		@GeneratedValue
+		private Long id;
+	}
+
+	@Entity(name = "Dog")
+	public static abstract class Dog extends Animal {
+		private boolean fast;
+
+		protected Dog(boolean fast) {
+			this.fast = fast;
+		}
+
+		public final boolean isFast() {
+			return fast;
+		}
+	}
+
+	@Entity(name = "Dachshund")
+	public static class Dachshund extends Dog {
+		public Dachshund() {
+			super( false );
+		}
+	}
+
+	@Entity(name = "Greyhound")
+	public static class Greyhound extends Dog {
+		public Greyhound() {
+			super( true );
+		}
+	}
 }

hibernate-entitymanager/src/test/java/org/hibernate/jpa/test/criteria/TreatKeywordTest.java

@@ -6,8 +6,13 @@
  */
 package org.hibernate.jpa.test.criteria;
 
+import java.util.Arrays;
 import java.util.List;
+import javax.persistence.Entity;
 import javax.persistence.EntityManager;
+import javax.persistence.EntityTransaction;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.criteria.CriteriaQuery;
 import javax.persistence.criteria.Root;
@@ -22,6 +27,8 @@ import org.hibernate.testing.TestForIssue;
 import org.junit.Assert;
 import org.junit.Test;
 
+import static org.junit.Assert.assertEquals;
+
 /**
  * @author Steve Ebersole
  */
@@ -29,7 +36,10 @@ public class TreatKeywordTest extends BaseEntityManagerFunctionalTestCase {
 
 	@Override
 	protected Class<?>[] getAnnotatedClasses() {
-		return new Class[] { Animal.class, Elephant.class, Human.class, Thing.class, ThingWithQuantity.class };
+		return new Class[] {
+				Animal.class, Elephant.class, Human.class, Thing.class, ThingWithQuantity.class,
+				TreatAnimal.class, Dog.class, Dachshund.class, Greyhound.class
+		};
 	}
 
 	@Test
@@ -213,4 +223,68 @@ public class TreatKeywordTest extends BaseEntityManagerFunctionalTestCase {
 	}
 
 
+	@Test
+	@TestForIssue(jiraKey = "HHH-9411")
+	public void testTreatWithRestrictionOnAbstractClass() {
+		EntityManager em = getOrCreateEntityManager();
+		EntityTransaction entityTransaction = em.getTransaction();
+		entityTransaction.begin();
+
+		Greyhound greyhound = new Greyhound();
+		Dachshund dachshund = new Dachshund();
+		em.persist( greyhound );
+		em.persist( dachshund );
+
+		CriteriaBuilder cb = em.getCriteriaBuilder();
+		CriteriaQuery<TreatAnimal> criteriaQuery = cb.createQuery( TreatAnimal.class );
+
+		Root<TreatAnimal> animal = criteriaQuery.from( TreatAnimal.class );
+
+		Root<Dog> dog = cb.treat( animal, Dog.class );
+
+		// only fast dogs
+		criteriaQuery.where( cb.isTrue( dog.<Boolean>get( "fast" ) ) );
+
+		List<TreatAnimal> results = em.createQuery( criteriaQuery ).getResultList();
+
+		// we should only have a single Greyhound here, not slow long dogs!
+		assertEquals( Arrays.asList( greyhound ), results );
+
+		entityTransaction.commit();
+		em.close();
+	}
+
+	@Entity
+	public static abstract class TreatAnimal {
+		@Id
+		@GeneratedValue
+		private Long id;
+	}
+
+	@Entity
+	public static abstract class Dog extends TreatAnimal {
+		private boolean fast;
+
+		protected Dog(boolean fast) {
+			this.fast = fast;
+		}
+
+		public final boolean isFast() {
+			return fast;
+		}
+	}
+
+	@Entity
+	public static class Dachshund extends Dog {
+		public Dachshund() {
+			super( false );
+		}
+	}
+
+	@Entity
+	public static class Greyhound extends Dog {
+		public Greyhound() {
+			super( true );
+		}
+	}
 }