HHH-19364 make CriteriaDefinition work nicely with Augmentation

Author: gavinking

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

@@ -169,6 +169,11 @@ public CriteriaDefinition(SessionFactory factory, CriteriaQuery<R> baseQuery) {
 		query = (JpaCriteriaQuery<R>) baseQuery;
 	}
 
+	public CriteriaDefinition(CriteriaQuery<R> baseQuery) {
+		super( ((JpaCriteriaQuery<R>) baseQuery).getCriteriaBuilder() );
+		query = (JpaCriteriaQuery<R>) baseQuery;
+	}
+
 	public CriteriaDefinition(EntityManagerFactory factory, Class<R> resultType) {
 		super( factory.getCriteriaBuilder() );
 		query = createQuery( resultType );

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

@@ -8,6 +8,7 @@
 import org.hibernate.query.IllegalMutationQueryException;
 import org.hibernate.query.IllegalSelectQueryException;
 import org.hibernate.query.Order;
+import org.hibernate.query.criteria.CriteriaDefinition;
 import org.hibernate.query.programmatic.MutationSpecification;
 import org.hibernate.query.programmatic.SelectionSpecification;
 import org.hibernate.query.range.Range;
@@ -195,7 +196,7 @@ void testCriteriaForm(SessionFactoryScope factoryScope) {
 	}
 
 	@Test
-	void testCriteriaFormWithMutation(SessionFactoryScope factoryScope) {
+	void testAugmentation(SessionFactoryScope factoryScope) {
 		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();
 
 		factoryScope.inTransaction( (session) -> {
@@ -211,6 +212,27 @@ void testCriteriaFormWithMutation(SessionFactoryScope factoryScope) {
 		assertThat( sqlCollector.getSqlQueries().get( 0 ) ).contains( " order by be1_0.position" );
 	}
 
+	@Test
+	void testAugmentationViaCriteriaDefinition(SessionFactoryScope factoryScope) {
+		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();
+
+		factoryScope.inTransaction( (session) -> {
+			sqlCollector.clear();
+			SelectionSpecification.create( BasicEntity.class )
+					.addAugmentation( (builder, query, entity) ->
+							new CriteriaDefinition<>( query ) {{
+								where( like( entity.get( BasicEntity_.name ), "%" ),
+										equal( entity.get( BasicEntity_.position), 1 ) );
+							}}
+					)
+					.addOrdering( Order.asc( BasicEntity_.position ) )
+					.createQuery( session )
+					.getResultList();
+		} );
+
+		assertThat( sqlCollector.getSqlQueries() ).hasSize( 1 );
+		assertThat( sqlCollector.getSqlQueries().get( 0 ) ).contains( " order by be1_0.position" );
+	}
 	@Test
 	void testBaseParameters(SessionFactoryScope factoryScope) {
 		final SQLStatementInspector sqlCollector = factoryScope.getCollectingStatementInspector();