HHH-18979 add Restriction.apply() so it can be easily used with criteria queries

Author: gavinking

hibernate-core/src/main/java/org/hibernate/query/criteria/CriteriaDefinition.java

@@ -148,8 +148,10 @@ public CriteriaDefinition(CriteriaDefinition<R> template) {
 	 */
 	public CriteriaDefinition(CriteriaDefinition<?> template, Class<R> resultType) {
 		super( template.getCriteriaBuilder() );
-		query = ((SqmSelectStatement<?>) template.query)
-				.createCopy( SqmCopyContext.simpleContext(), resultType );
+		if ( !(template.query instanceof SqmSelectStatement<?> selectStatement) ) {
+			throw new IllegalArgumentException( "Not a SqmSelectStatement" );
+		}
+		query = selectStatement.createCopy( SqmCopyContext.simpleContext(), resultType );
 	}
 
 	public CriteriaDefinition(SessionFactory factory, Class<R> resultType) {
@@ -220,6 +222,11 @@ public JpaCriteriaQuery<R> restrict(Predicate predicate) {
 		return existing == null ? where( predicate ) : where( existing, predicate );
 	}
 
+	@Override
+	public HibernateCriteriaBuilder getCriteriaBuilder() {
+		return query.getCriteriaBuilder();
+	}
+
 	@Override
 	public JpaCriteriaQuery<R> select(Selection<? extends R> selection) {
 		return query.select(selection);

hibernate-core/src/main/java/org/hibernate/query/criteria/JpaCriteriaQuery.java

@@ -132,4 +132,6 @@ default List<Order> getOrderList() {
 
 	@Override
 	<U> JpaSubQuery<U> subquery(EntityType<U> type);
+
+	HibernateCriteriaBuilder getCriteriaBuilder();
 }

hibernate-core/src/main/java/org/hibernate/query/restriction/Restriction.java

@@ -5,13 +5,15 @@
 package org.hibernate.query.restriction;
 
 import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.Predicate;
 import jakarta.persistence.criteria.Root;
 import jakarta.persistence.metamodel.SingularAttribute;
 import org.hibernate.Incubating;
 import org.hibernate.Internal;
 import org.hibernate.query.Order;
 import org.hibernate.query.SelectionQuery;
+import org.hibernate.query.criteria.JpaCriteriaQuery;
 import org.hibernate.query.range.Range;
 
 import java.util.List;
@@ -77,6 +79,17 @@ default Restriction<X> and(Restriction<X> restriction) {
 	@Internal
 	Predicate toPredicate(Root<? extends X> root, CriteriaBuilder builder);
 
+	/**
+	 * Apply this restriction to the given root entity of the given
+	 * {@linkplain CriteriaQuery criteria query}.
+	 */
+	default void apply(CriteriaQuery<?> query, Root<? extends X> root) {
+		if ( !(query instanceof JpaCriteriaQuery<?> criteriaQuery) ) {
+			throw new IllegalArgumentException( "Not a JpaCriteriaQuery" );
+		}
+		query.where( query.getRestriction(), toPredicate( root, criteriaQuery.getCriteriaBuilder() ) );
+	}
+
 	/**
 	 * Restrict the allowed values of the given attribute to the given
 	 * {@linkplain Range range}.

hibernate-core/src/main/java/org/hibernate/query/sqm/internal/AbstractSqmSelectionQuery.java

@@ -150,9 +150,13 @@ public SelectionQuery<R> setOrder(Order<? super R> order) {
 	@Override
 	public SelectionQuery<R> addRestriction(Restriction<? super R> restriction) {
 		final SqmSelectStatement<R> selectStatement = getSqmSelectStatement().copy( noParamCopyContext() );
-		final Root<? extends R> root = (Root<? extends R>) selectStatement.getRootList().get( 0 );
-		selectStatement.where( selectStatement.getRestriction(),
-				restriction.toPredicate( root, selectStatement.nodeBuilder() ) );
+		final Root<?> firstRoot = selectStatement.getRootList().get( 0 );
+		if ( !getExpectedResultType().isAssignableFrom( firstRoot.getJavaType() ) ) {
+			throw new IllegalStateException("First root entity of the query did not have the query result type");
+		}
+		@SuppressWarnings("unchecked") // safe, we just checked
+		final Root<? extends R> root = (Root<? extends R>) firstRoot;
+		restriction.apply( selectStatement, root );
 		// TODO: when the QueryInterpretationCache can handle caching criteria queries,
 		//       simply cache the new SQM as if it were a criteria query, and remove this:
 		getQueryOptions().setQueryPlanCachingEnabled( false );

hibernate-core/src/main/java/org/hibernate/query/sqm/tree/select/SqmSelectStatement.java

@@ -20,6 +20,7 @@
 import org.hibernate.query.sqm.NodeBuilder;
 import org.hibernate.query.sqm.SemanticQueryWalker;
 import org.hibernate.query.sqm.SqmQuerySource;
+import org.hibernate.query.sqm.internal.ParameterCollector;
 import org.hibernate.query.sqm.internal.SqmUtil;
 import org.hibernate.query.sqm.tree.SqmCopyContext;
 import org.hibernate.query.sqm.tree.SqmStatement;
@@ -49,8 +50,8 @@
 /**
  * @author Steve Ebersole
  */
-public class SqmSelectStatement<T> extends AbstractSqmSelectQuery<T> implements JpaCriteriaQuery<T>, SqmStatement<T>,
-		org.hibernate.query.sqm.internal.ParameterCollector {
+public class SqmSelectStatement<T> extends AbstractSqmSelectQuery<T>
+		implements JpaCriteriaQuery<T>, SqmStatement<T>, ParameterCollector {
 	private final SqmQuerySource querySource;
 
 	private Set<SqmParameter<?>> parameters;
@@ -160,6 +161,11 @@ public void validateResultType(Class<?> resultType) {
 		SqmUtil.validateQueryReturnType( getQueryPart(), resultType );
 	}
 
+	@Override
+	public NodeBuilder getCriteriaBuilder() {
+		return nodeBuilder();
+	}
+
 	@Override
 	public SqmQuerySource getQuerySource() {
 		return querySource;

hibernate-core/src/test/java/org/hibernate/orm/test/query/restriction/RestrictionTest.java

@@ -8,6 +8,7 @@
 import jakarta.persistence.Id;
 import jakarta.persistence.ManyToOne;
 import jakarta.persistence.Version;
+import jakarta.persistence.criteria.Root;
 import jakarta.persistence.metamodel.SingularAttribute;
 import org.hibernate.testing.orm.junit.DomainModel;
 import org.hibernate.testing.orm.junit.SessionFactory;
@@ -216,6 +217,48 @@ void testPath(SessionFactoryScope scope) {
 		assertEquals( 2, unsafeTest.size() );
 	}
 
+	@Test
+	void testCriteria(SessionFactoryScope scope) {
+		scope.getSessionFactory().getSchemaManager().truncate();
+		scope.inTransaction( session -> {
+			session.persist(new Book("9781932394153", "Hibernate in Action", 400));
+			session.persist(new Book("9781617290459", "Java Persistence with Hibernate", 1000));
+		});
+
+		var bookType = scope.getSessionFactory().getJpaMetamodel().findEntityType(Book.class);
+		@SuppressWarnings( "unchecked" )
+		var title = (SingularAttribute<? super Book, String>) bookType.findSingularAttribute("title");
+		@SuppressWarnings( "unchecked" )
+		var isbn = (SingularAttribute<? super Book, String>) bookType.findSingularAttribute("isbn");
+		@SuppressWarnings( "unchecked" )
+		var pages = (SingularAttribute<? super Book, Integer>) bookType.findSingularAttribute("pages");
+
+		scope.inSession( session -> {
+			var query = session.getCriteriaBuilder().createQuery(String.class);
+			var root = query.from( Book.class );
+			like( title, "%Hibernate%" ).apply( query, root );
+			query.select( root.get( title ) );
+			List<String> titles = session.createQuery( query ).getResultList();
+			assertEquals( 2, titles.size() );
+		} );
+		scope.inSession( session -> {
+			var query = session.getCriteriaBuilder().createQuery(String.class);
+			var root = query.from( Book.class );
+			equal( isbn, "9781932394153" ).apply( query, root );
+			query.select( root.get( title ) );
+			List<String> titles = session.createQuery( query ).getResultList();
+			assertEquals( 1, titles.size() );
+		} );
+		scope.inSession( session -> {
+			var query = session.getCriteriaBuilder().createQuery("select title from Book", String.class);
+			var root = (Root<Book>) query.getRootList().get(0);
+			equal( isbn, "9781932394153" ).apply( query, root );
+			List<String> titles = session.createQuery( query ).getResultList();
+			assertEquals( 1, titles.size() );
+		} );
+	}
+
+
 	@Entity(name="Book")
 	static class Book {
 		@Id