HHH-19364 make Specification subtypes into TypedQueryReferences

Author: gavinking

hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java

@@ -10,6 +10,8 @@
 import jakarta.persistence.criteria.CriteriaDelete;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.CriteriaUpdate;
+import jakarta.persistence.criteria.CommonAbstractCriteria;
+
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.hibernate.CacheMode;
 import org.hibernate.EntityNameResolver;
@@ -71,6 +73,8 @@
 import org.hibernate.query.hql.spi.SqmQueryImplementor;
 import org.hibernate.query.named.NamedObjectRepository;
 import org.hibernate.query.named.NamedResultSetMappingMemento;
+import org.hibernate.query.programmatic.MutationSpecification;
+import org.hibernate.query.programmatic.SelectionSpecification;
 import org.hibernate.query.spi.HqlInterpretation;
 import org.hibernate.query.spi.QueryImplementor;
 import org.hibernate.query.sql.internal.NativeQueryImpl;
@@ -901,15 +905,25 @@ public <T> QueryImplementor<T> createQuery(String queryString, Class<T> expected
 	@Override
 	public <R> QueryImplementor<R> createQuery(TypedQueryReference<R> typedQueryReference) {
 		checksBeforeQueryCreation();
-		@SuppressWarnings("unchecked")
-		// this cast is fine because of all our impls of TypedQueryReference return Class<R>
-		final Class<R> resultType = (Class<R>) typedQueryReference.getResultType();
-		final QueryImplementor<R> query =
-				buildNamedQuery( typedQueryReference.getName(),
-						memento -> createSqmQueryImplementor( resultType, memento ),
-						memento -> createNativeQueryImplementor( resultType, memento ) );
-		typedQueryReference.getHints().forEach( query::setHint );
-		return query;
+		if ( typedQueryReference instanceof SelectionSpecification<R> specification ) {
+			final CriteriaQuery<R> query = specification.buildCriteria( getCriteriaBuilder() );
+			return new QuerySqmImpl<>( (SqmStatement<R>) query, specification.getResultType(), this );
+		}
+		else if ( typedQueryReference instanceof MutationSpecification<?> specification ) {
+			final CommonAbstractCriteria query = specification.buildCriteria( getCriteriaBuilder() );
+			return new QuerySqmImpl<>( (SqmStatement<R>) query, (Class<R>) specification.getResultType(), this );
+		}
+		else {
+			@SuppressWarnings("unchecked")
+			// this cast is fine because of all our impls of TypedQueryReference return Class<R>
+			final Class<R> resultType = (Class<R>) typedQueryReference.getResultType();
+			final QueryImplementor<R> query =
+					buildNamedQuery( typedQueryReference.getName(),
+							memento -> createSqmQueryImplementor( resultType, memento ),
+							memento -> createNativeQueryImplementor( resultType, memento ) );
+			typedQueryReference.getHints().forEach( query::setHint );
+			return query;
+		}
 	}
 
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

hibernate-core/src/main/java/org/hibernate/query/programmatic/MutationSpecification.java

@@ -4,11 +4,13 @@
  */
 package org.hibernate.query.programmatic;
 
+import jakarta.persistence.TypedQueryReference;
 import jakarta.persistence.criteria.CommonAbstractCriteria;
 import jakarta.persistence.criteria.CriteriaUpdate;
 import jakarta.persistence.criteria.CriteriaDelete;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.Root;
+
 import org.hibernate.Incubating;
 import org.hibernate.Session;
 import org.hibernate.StatelessSession;
@@ -37,7 +39,7 @@
  * @since 7.0
  */
 @Incubating
-public interface MutationSpecification<T> extends QuerySpecification<T> {
+public interface MutationSpecification<T> extends QuerySpecification<T>, TypedQueryReference<Void> {
 
 	/**
 	 * Covariant override.
@@ -79,6 +81,9 @@ interface Augmentation<T> {
 	@Override
 	MutationSpecification<T> validate(CriteriaBuilder builder);
 
+	@Override
+	Class<Void> getResultType();
+
 	/**
 	 * Returns a specification reference which can be used to programmatically,
 	 * iteratively build a {@linkplain MutationQuery} based on a base HQL statement,

hibernate-core/src/main/java/org/hibernate/query/programmatic/SelectionSpecification.java

@@ -5,6 +5,7 @@
 package org.hibernate.query.programmatic;
 
 import jakarta.persistence.EntityManager;
+import jakarta.persistence.TypedQueryReference;
 import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.CriteriaBuilder;
 import jakarta.persistence.criteria.Root;
@@ -59,7 +60,7 @@
  * @since 7.0
  */
 @Incubating
-public interface SelectionSpecification<T> extends QuerySpecification<T> {
+public interface SelectionSpecification<T> extends QuerySpecification<T>, TypedQueryReference<T> {
 	/**
 	 * Adds an ordering to the selection specification.
 	 * Appended to any previous ordering.
@@ -185,6 +186,9 @@ interface Augmentation<T> {
 	@Override
 	SelectionSpecification<T> validate(CriteriaBuilder builder);
 
+	@Override
+	Class<T> getResultType();
+
 	/**
 	 * Returns a specification reference which can be used to programmatically,
 	 * iteratively build a {@linkplain SelectionQuery} for the given entity type,

hibernate-core/src/main/java/org/hibernate/query/programmatic/internal/MutationSpecificationImpl.java

@@ -33,8 +33,10 @@
 import org.hibernate.query.sqm.tree.update.SqmUpdateStatement;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.function.BiConsumer;
 
 import static org.hibernate.query.sqm.tree.SqmCopyContext.noParamCopyContext;
@@ -69,6 +71,21 @@ public MutationSpecificationImpl(CriteriaDelete<T> criteriaQuery) {
 		this.hql = null;
 	}
 
+	@Override
+	public String getName() {
+		return null;
+	}
+
+	@Override
+	public Class<Void> getResultType() {
+		return Void.class;
+	}
+
+	@Override
+	public Map<String,Object> getHints() {
+		return Collections.emptyMap();
+	}
+
 	@Override
 	public MutationSpecification<T> restrict(Restriction<? super T> restriction) {
 		specifications.add( (sqmStatement, mutationTargetRoot) -> {

hibernate-core/src/main/java/org/hibernate/query/programmatic/internal/SelectionSpecificationImpl.java

@@ -32,6 +32,7 @@
 import org.hibernate.query.sqm.tree.select.SqmSortSpecification;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -70,6 +71,21 @@ public SelectionSpecificationImpl(CriteriaQuery<T> criteriaQuery) {
 		this.criteriaQuery = criteriaQuery;
 	}
 
+	@Override
+	public String getName() {
+		return null;
+	}
+
+	@Override
+	public Class<T> getResultType() {
+		return resultType;
+	}
+
+	@Override
+	public Map<String,Object> getHints() {
+		return Collections.emptyMap();
+	}
+
 	@Override
 	public SelectionSpecification<T> restrict(Restriction<? super T> restriction) {
 		specifications.add( (sqmStatement, root) -> {

hibernate-core/src/test/java/org/hibernate/orm/test/query/dynamic/SimpleQuerySpecificationTests.java

@@ -296,4 +296,21 @@ void testBuildCriteriaQuery(SessionFactoryScope factoryScope) {
 						.validate( factory.getCriteriaBuilder() )
 						.buildCriteria( factory.getCriteriaBuilder() );
 	}
+
+	@Test
+	void testUseAsTypedQueryRef(SessionFactoryScope factoryScope) {
+		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();
+
+		factoryScope.inTransaction( (session) -> {
+			sqlCollector.clear();
+			var spec = SelectionSpecification.create( BasicEntity.class )
+					.restrict( Restriction.restrict( BasicEntity_.position, Range.closed( 1, 5 ) ) )
+					.sort( Order.asc( BasicEntity_.position ) );
+			session.createQuery(spec).getResultList();
+		} );
+
+		assertThat( sqlCollector.getSqlQueries() ).hasSize( 1 );
+		assertThat( sqlCollector.getSqlQueries().get( 0 ) ).contains( " order by be1_0.position" );
+	}
+
 }