Skip to content

Commit 81f983a

Browse files
DavideDDavideD
committed
[#1932] Support ReactiveSelectionQuery#getResultCount for HQL
1 parent 46291aa commit 81f983a

14 files changed

+203
-24
lines changed

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/Mutiny.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,17 @@ interface SelectionQuery<R> extends AbstractQuery {
194194
*/
195195
Uni<R> getSingleResultOrNull();
196196

197+
/**
198+
* Determine the size of the query result list that would be
199+
* returned by calling {@link #getResultList()} with no
200+
* {@linkplain #getFirstResult() offset} or
201+
* {@linkplain #getMaxResults() limit} applied to the query.
202+
*
203+
* @return the size of the list that would be returned
204+
*/
205+
@Incubating
206+
Uni<Long> getResultCount();
207+
197208
/**
198209
* Asynchronously execute this query, returning the query results
199210
* as a {@link List}, via a {@link Uni}. If the query

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinyQueryImpl.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ public int getMaxResults() {
4747
return delegate.getMaxResults();
4848
}
4949

50+
@Override
51+
public Uni<Long> getResultCount() {
52+
return uni( delegate::getReactiveResultCount );
53+
}
54+
5055
@Override
5156
public Uni<List<R>> getResultList() {
5257
return uni( delegate::getReactiveResultList );

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySelectionQueryImpl.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,11 @@ public int getMaxResults() {
4545
return delegate.getMaxResults();
4646
}
4747

48+
@Override
49+
public Uni<Long> getResultCount() {
50+
return uni( delegate::getReactiveResultCount );
51+
}
52+
4853
@Override
4954
public Uni<List<R>> getResultList() {
5055
return uni( delegate::getReactiveResultList );

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/ReactiveSelectionQuery.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ default CompletionStage<List<R>> getReactiveResultList() {
4949

5050
CompletionStage<R> getReactiveSingleResultOrNull();
5151

52+
CompletionStage<Long> getReactiveResultCount();
53+
5254
CompletionStage<R> reactiveUnique();
5355

5456
CompletionStage<Optional<R>> reactiveUniqueResultOptional();

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/spi/ReactiveAbstractSelectionQuery.java

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import org.hibernate.engine.spi.SharedSessionContractImplementor;
2222
import org.hibernate.query.IllegalQueryOperationException;
2323
import org.hibernate.query.hql.internal.QuerySplitter;
24+
import org.hibernate.query.internal.DelegatingDomainQueryExecutionContext;
25+
import org.hibernate.query.spi.DomainQueryExecutionContext;
2426
import org.hibernate.query.spi.QueryInterpretationCache;
2527
import org.hibernate.query.spi.QueryOptions;
2628
import org.hibernate.query.sqm.internal.DomainParameterXref;
@@ -33,6 +35,7 @@
3335
import org.hibernate.reactive.query.sqm.internal.AggregatedSelectReactiveQueryPlan;
3436
import org.hibernate.reactive.query.sqm.internal.ConcreteSqmSelectReactiveQueryPlan;
3537
import org.hibernate.reactive.query.sqm.spi.ReactiveSelectQueryPlan;
38+
import org.hibernate.reactive.sql.results.spi.ReactiveSingleResultConsumer;
3639
import org.hibernate.sql.results.internal.TupleMetadata;
3740

3841
import jakarta.persistence.NoResultException;
@@ -146,6 +149,17 @@ public CompletionStage<R> getReactiveSingleResult() {
146149
.exceptionally( this::convertException );
147150
}
148151

152+
public CompletionStage<Long> getReactiveResultsCount(SqmSelectStatement<?> sqmStatement, DomainQueryExecutionContext domainQueryExecutionContext) {
153+
final DelegatingDomainQueryExecutionContext context = new DelegatingDomainQueryExecutionContext( domainQueryExecutionContext ) {
154+
@Override
155+
public QueryOptions getQueryOptions() {
156+
return QueryOptions.NONE;
157+
}
158+
};
159+
return buildConcreteSelectQueryPlan( sqmStatement.createCountQuery(), Long.class, getQueryOptions() )
160+
.reactiveExecuteQuery( context, new ReactiveSingleResultConsumer<>() );
161+
}
162+
149163
private R reactiveSingleResult(List<R> list) {
150164
if ( list.isEmpty() ) {
151165
throw new NoResultException( String.format( "No result found for query [%s]", getQueryString() ) );
@@ -269,7 +283,7 @@ private ReactiveSelectQueryPlan<R> buildAggregatedSelectQueryPlan(SqmSelectState
269283
return new AggregatedSelectReactiveQueryPlan<>( aggregatedQueryPlans );
270284
}
271285

272-
private <T> ReactiveSelectQueryPlan<T> buildConcreteSelectQueryPlan(
286+
public <T> ReactiveSelectQueryPlan<T> buildConcreteSelectQueryPlan(
273287
SqmSelectStatement<?> concreteSqmStatement,
274288
Class<T> resultType,
275289
QueryOptions queryOptions) {

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/internal/ConcreteSqmSelectReactiveQueryPlan.java

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.hibernate.reactive.query.sqm.spi.ReactiveSelectQueryPlan;
3737
import org.hibernate.reactive.sql.exec.internal.StandardReactiveSelectExecutor;
3838
import org.hibernate.reactive.sql.results.spi.ReactiveListResultsConsumer;
39+
import org.hibernate.reactive.sql.results.spi.ReactiveResultsConsumer;
3940
import org.hibernate.sql.ast.SqlAstTranslator;
4041
import org.hibernate.sql.ast.SqlAstTranslatorFactory;
4142
import org.hibernate.sql.ast.spi.FromClauseAccess;
@@ -59,6 +60,7 @@
5960
public class ConcreteSqmSelectReactiveQueryPlan<R> extends ConcreteSqmSelectQueryPlan<R>
6061
implements ReactiveSelectQueryPlan<R> {
6162

63+
private final SqmInterpreter<Object, ReactiveResultsConsumer<Object, R>> executeQueryInterpreter;
6264
private final SqmInterpreter<List<R>, Void> listInterpreter;
6365
private final RowTransformer<R> rowTransformer;
6466

@@ -80,6 +82,8 @@ public ConcreteSqmSelectReactiveQueryPlan(
8082
this.rowTransformer = determineRowTransformer( sqm, resultType, tupleMetadata, queryOptions );
8183
this.listInterpreter = (unused, executionContext, sqmInterpretation, jdbcParameterBindings) ->
8284
listInterpreter( hql, domainParameterXref, executionContext, sqmInterpretation, jdbcParameterBindings, rowTransformer );
85+
this.executeQueryInterpreter = (resultsConsumer, executionContext, sqmInterpretation, jdbcParameterBindings) ->
86+
executeQueryInterpreter( hql, domainParameterXref, executionContext, sqmInterpretation, jdbcParameterBindings, rowTransformer, resultsConsumer );
8387
}
8488

8589
private static <R> CompletionStage<List<R>> listInterpreter(
@@ -110,6 +114,40 @@ private static <R> CompletionStage<List<R>> listInterpreter(
110114
.whenComplete( (rs, t) -> domainParameterXref.clearExpansions() );
111115
}
112116

117+
private static <R> CompletionStage<Object> executeQueryInterpreter(
118+
String hql,
119+
DomainParameterXref domainParameterXref,
120+
DomainQueryExecutionContext executionContext,
121+
CacheableSqmInterpretation sqmInterpretation,
122+
JdbcParameterBindings jdbcParameterBindings,
123+
RowTransformer<R> rowTransformer,
124+
ReactiveResultsConsumer<Object, R> resultsConsumer) {
125+
final ReactiveSharedSessionContractImplementor session = (ReactiveSharedSessionContractImplementor) executionContext.getSession();
126+
final JdbcOperationQuerySelect jdbcSelect = sqmInterpretation.getJdbcSelect();
127+
// I'm using a supplier so that the whenComplete at the end will catch any errors, like a finally-block
128+
Supplier<SubselectFetch.RegistrationHandler> fetchHandlerSupplier = () -> SubselectFetch
129+
.createRegistrationHandler( session.getPersistenceContext().getBatchFetchQueue(), sqmInterpretation.selectStatement, JdbcParametersList.empty(), jdbcParameterBindings );
130+
return completedFuture( fetchHandlerSupplier )
131+
.thenApply( Supplier::get )
132+
.thenCompose( subSelectFetchKeyHandler -> session
133+
.reactiveAutoFlushIfRequired( jdbcSelect.getAffectedTableNames() )
134+
.thenCompose( required -> StandardReactiveSelectExecutor.INSTANCE
135+
.executeQuery( jdbcSelect,
136+
jdbcParameterBindings,
137+
ConcreteSqmSelectQueryPlan.listInterpreterExecutionContext( hql, executionContext, jdbcSelect, subSelectFetchKeyHandler ),
138+
rowTransformer,
139+
null,
140+
sql -> executionContext.getSession()
141+
.getJdbcCoordinator()
142+
.getStatementPreparer()
143+
.prepareQueryStatement( sql, false, null ),
144+
resultsConsumer
145+
)
146+
)
147+
)
148+
.whenComplete( (rs, t) -> domainParameterXref.clearExpansions() );
149+
}
150+
113151
@Override
114152
public ScrollableResultsImplementor<R> performScroll(ScrollMode scrollMode, DomainQueryExecutionContext executionContext) {
115153
throw new UnsupportedOperationException();
@@ -119,10 +157,21 @@ public ScrollableResultsImplementor<R> performScroll(ScrollMode scrollMode, Doma
119157
public CompletionStage<List<R>> reactivePerformList(DomainQueryExecutionContext executionContext) {
120158
return executionContext.getQueryOptions().getEffectiveLimit().getMaxRowsJpa() == 0
121159
? completedFuture( emptyList() )
122-
: withCacheableSqmInterpretation( executionContext, listInterpreter );
160+
: withCacheableSqmInterpretation( executionContext, null, listInterpreter );
161+
}
162+
163+
@Override
164+
public <T> CompletionStage<T> reactiveExecuteQuery(
165+
DomainQueryExecutionContext executionContext,
166+
ReactiveResultsConsumer<T, R> resultsConsumer) {
167+
return withCacheableSqmInterpretation(
168+
executionContext,
169+
resultsConsumer,
170+
(SqmInterpreter<T, ReactiveResultsConsumer<T, R>>) (SqmInterpreter) executeQueryInterpreter
171+
);
123172
}
124173

125-
private <T, X> CompletionStage<T> withCacheableSqmInterpretation(DomainQueryExecutionContext executionContext, SqmInterpreter<T, X> interpreter) {
174+
private <T, X> CompletionStage<T> withCacheableSqmInterpretation(DomainQueryExecutionContext executionContext, X context, SqmInterpreter<T, X> interpreter) {
126175
// NOTE : VERY IMPORTANT - intentional double-lock checking
127176
// The other option would be to leverage `java.util.concurrent.locks.ReadWriteLock`
128177
// to protect access. However, synchronized is much simpler here. We will verify
@@ -162,7 +211,7 @@ private <T, X> CompletionStage<T> withCacheableSqmInterpretation(DomainQueryExec
162211
jdbcParameterBindings = createJdbcParameterBindings( localCopy, executionContext );
163212
}
164213

165-
return interpreter.interpret( null, executionContext, localCopy, jdbcParameterBindings );
214+
return interpreter.interpret( context, executionContext, localCopy, jdbcParameterBindings );
166215
}
167216

168217
private JdbcParameterBindings createJdbcParameterBindings(CacheableSqmInterpretation sqmInterpretation, DomainQueryExecutionContext executionContext) {

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/internal/ReactiveQuerySqmImpl.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,12 @@ public CompletionStage<R> getReactiveSingleResult() {
139139
return selectionQueryDelegate.getReactiveSingleResult();
140140
}
141141

142+
@Override
143+
public CompletionStage<Long> getReactiveResultCount() {
144+
return selectionQueryDelegate
145+
.getReactiveResultsCount( ( (SqmSelectStatement<?>) getSqmStatement() ).createCountQuery(), this );
146+
}
147+
142148
@Override
143149
public CompletionStage<R> getReactiveSingleResultOrNull() {
144150
return selectionQueryDelegate.getReactiveSingleResultOrNull();

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/internal/ReactiveSqmSelectionQueryImpl.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,12 @@ public R getSingleResultOrNull() {
212212
return selectionQueryDelegate.getSingleResultOrNull();
213213
}
214214

215+
@Override
216+
public CompletionStage<Long> getReactiveResultCount() {
217+
return selectionQueryDelegate
218+
.getReactiveResultsCount( getSqmStatement().createCountQuery(), this );
219+
}
220+
215221
@Override
216222
public List<R> getResultList() {
217223
return selectionQueryDelegate.getResultList();

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/spi/ReactiveSelectQueryPlan.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.hibernate.query.spi.ScrollableResultsImplementor;
1515
import org.hibernate.query.spi.SelectQueryPlan;
1616
import org.hibernate.reactive.logging.impl.Log;
17+
import org.hibernate.reactive.sql.results.spi.ReactiveResultsConsumer;
1718
import org.hibernate.sql.results.spi.ResultsConsumer;
1819

1920
import static org.hibernate.reactive.logging.impl.LoggerFactory.make;
@@ -44,7 +45,7 @@ default <T> T executeQuery(DomainQueryExecutionContext executionContext, Results
4445
/**
4546
* Execute the query
4647
*/
47-
default <T> CompletionStage<T> reactiveExecuteQuery(DomainQueryExecutionContext executionContext, ResultsConsumer<T, R> resultsConsumer) {
48+
default <T> CompletionStage<T> reactiveExecuteQuery(DomainQueryExecutionContext executionContext, ReactiveResultsConsumer<T, R> resultsConsumer) {
4849
return failedFuture( new UnsupportedOperationException() );
4950
}
5051

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/sql/results/internal/ReactiveInitializersList.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,12 @@ public void finishUpRow(final RowProcessingState rowProcessingState) {
5656
}
5757
}
5858

59+
public void startLoading(final RowProcessingState rowProcessingState) {
60+
for ( int i = initializers.length - 1; i >= 0; i-- ) {
61+
initializers[i].startLoading( rowProcessingState );
62+
}
63+
}
64+
5965
public CompletionStage<Void> initializeInstance(final ReactiveRowProcessingState rowProcessingState) {
6066
return loop( initializers, initializer -> {
6167
if ( initializer instanceof ReactiveInitializer ) {

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/sql/results/spi/ReactiveSingleResultConsumer.java

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/* Hibernate, Relational Persistence for Idiomatic Java
2+
*
3+
* SPDX-License-Identifier: Apache-2.0
4+
* Copyright: Red Hat Inc. and Hibernate Authors
5+
*/
6+
package org.hibernate.reactive.sql.results.spi;
7+
8+
import java.util.concurrent.CompletionStage;
9+
10+
import org.hibernate.Incubating;
11+
import org.hibernate.engine.spi.SharedSessionContractImplementor;
12+
import org.hibernate.reactive.sql.exec.spi.ReactiveRowProcessingState;
13+
import org.hibernate.reactive.sql.exec.spi.ReactiveValuesResultSet;
14+
import org.hibernate.sql.results.jdbc.internal.JdbcValuesSourceProcessingStateStandardImpl;
15+
import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingOptions;
16+
17+
@Incubating
18+
public class ReactiveSingleResultConsumer<T> implements ReactiveResultsConsumer<T, T> {
19+
20+
@Override
21+
public CompletionStage<T> consume(
22+
ReactiveValuesResultSet jdbcValues,
23+
SharedSessionContractImplementor session,
24+
JdbcValuesSourceProcessingOptions processingOptions,
25+
JdbcValuesSourceProcessingStateStandardImpl jdbcValuesSourceProcessingState,
26+
ReactiveRowProcessingState rowProcessingState,
27+
ReactiveRowReader<T> rowReader) {
28+
rowReader.getReactiveInitializersList().startLoading( rowProcessingState );
29+
return rowProcessingState.next()
30+
.thenCompose( hasNext -> rowReader
31+
.reactiveReadRow( rowProcessingState, processingOptions )
32+
.thenApply( result -> {
33+
rowProcessingState.finishRowProcessing( true );
34+
rowReader.finishUp( jdbcValuesSourceProcessingState );
35+
jdbcValuesSourceProcessingState.finishUp( false );
36+
return result;
37+
} )
38+
);
39+
}
40+
41+
@Override
42+
public boolean canResultsBeCached() {
43+
return false;
44+
}
45+
46+
}

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/stage/Stage.java

Lines changed: 37 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,20 +5,18 @@
55
*/
66
package org.hibernate.reactive.stage;
77

8-
import jakarta.persistence.CacheRetrieveMode;
9-
import jakarta.persistence.CacheStoreMode;
10-
import jakarta.persistence.EntityGraph;
11-
import jakarta.persistence.FlushModeType;
12-
import jakarta.persistence.LockModeType;
13-
import jakarta.persistence.Parameter;
14-
import jakarta.persistence.criteria.CriteriaBuilder;
15-
import jakarta.persistence.criteria.CriteriaDelete;
16-
import jakarta.persistence.criteria.CriteriaQuery;
17-
import jakarta.persistence.criteria.CriteriaUpdate;
18-
import jakarta.persistence.metamodel.Attribute;
19-
import jakarta.persistence.metamodel.Metamodel;
8+
import java.lang.invoke.MethodHandles;
9+
import java.util.List;
10+
import java.util.concurrent.CompletionStage;
11+
import java.util.function.BiFunction;
12+
import java.util.function.Function;
13+
2014
import org.hibernate.Cache;
21-
import org.hibernate.*;
15+
import org.hibernate.CacheMode;
16+
import org.hibernate.Filter;
17+
import org.hibernate.FlushMode;
18+
import org.hibernate.Incubating;
19+
import org.hibernate.LockMode;
2220
import org.hibernate.bytecode.enhance.spi.interceptor.EnhancementAsProxyLazinessInterceptor;
2321
import org.hibernate.collection.spi.AbstractPersistentCollection;
2422
import org.hibernate.collection.spi.PersistentCollection;
@@ -38,16 +36,25 @@
3836
import org.hibernate.reactive.util.impl.CompletionStages;
3937
import org.hibernate.stat.Statistics;
4038

41-
import java.lang.invoke.MethodHandles;
42-
import java.util.List;
43-
import java.util.concurrent.CompletionStage;
44-
import java.util.function.BiFunction;
45-
import java.util.function.Function;
39+
import jakarta.persistence.CacheRetrieveMode;
40+
import jakarta.persistence.CacheStoreMode;
41+
import jakarta.persistence.EntityGraph;
42+
import jakarta.persistence.FlushModeType;
43+
import jakarta.persistence.LockModeType;
44+
import jakarta.persistence.Parameter;
45+
import jakarta.persistence.criteria.CriteriaBuilder;
46+
import jakarta.persistence.criteria.CriteriaDelete;
47+
import jakarta.persistence.criteria.CriteriaQuery;
48+
import jakarta.persistence.criteria.CriteriaUpdate;
49+
import jakarta.persistence.metamodel.Attribute;
50+
import jakarta.persistence.metamodel.Metamodel;
4651

4752
import static org.hibernate.engine.internal.ManagedTypeHelper.asPersistentAttributeInterceptable;
4853
import static org.hibernate.engine.internal.ManagedTypeHelper.isPersistentAttributeInterceptable;
4954
import static org.hibernate.internal.util.LockModeConverter.convertToLockMode;
50-
import static org.hibernate.jpa.internal.util.CacheModeHelper.*;
55+
import static org.hibernate.jpa.internal.util.CacheModeHelper.interpretCacheMode;
56+
import static org.hibernate.jpa.internal.util.CacheModeHelper.interpretCacheRetrieveMode;
57+
import static org.hibernate.jpa.internal.util.CacheModeHelper.interpretCacheStoreMode;
5158

5259
/**
5360
* An API for Hibernate Reactive where non-blocking operations are
@@ -188,6 +195,17 @@ interface SelectionQuery<R> extends AbstractQuery {
188195
*/
189196
CompletionStage<R> getSingleResultOrNull();
190197

198+
/**
199+
* Determine the size of the query result list that would be
200+
* returned by calling {@link #getResultList()} with no
201+
* {@linkplain #getFirstResult() offset} or
202+
* {@linkplain #getMaxResults() limit} applied to the query.
203+
*
204+
* @return the size of the list that would be returned
205+
*/
206+
@Incubating
207+
CompletionStage<Long> getResultCount();
208+
191209
/**
192210
* Asynchronously execute this query, returning the query results
193211
* as a {@link List}, via a {@link CompletionStage}. If the query

‎hibernate-reactive-core/src/main/java/org/hibernate/reactive/stage/impl/StageQueryImpl.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@ public int getMaxResults() {
3838
return delegate.getMaxResults();
3939
}
4040

41+
@Override
42+
public CompletionStage<Long> getResultCount() {
43+
return delegate.getReactiveResultCount();
44+
}
45+
4146
@Override
4247
public CompletionStage<List<R>> getResultList() {
4348
return delegate.getReactiveResultList();

0 commit comments

Comments
 (0)