Skip to content

Commit ecc59ec

Browse files
DavideDDavideD
committed
[#1932] Support ReactiveSelectionQuery#getResultCount for HQL
I need some changes in ORM before I can support it for native queries.
1 parent f4ee6da commit ecc59ec

15 files changed

+218
-24
lines changed

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

+11
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

+5
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

+5
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

+2
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

+15-1
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/sql/internal/ReactiveNativeQueryImpl.java

+10
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,16 @@ private <T> T getNull() {
9393
return null;
9494
}
9595

96+
@Override
97+
public long getResultCount() {
98+
throw LOG.nonReactiveMethodCall( "getReactiveResultCount()" );
99+
}
100+
101+
@Override
102+
public CompletionStage<Long> getReactiveResultCount() {
103+
throw LOG.notYetImplemented();
104+
}
105+
96106
private ReactiveAbstractSelectionQuery<R> createSelectionQueryDelegate(SharedSessionContractImplementor session) {
97107
return new ReactiveAbstractSelectionQuery<>(
98108
this::getQueryOptions,

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

+52-3
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

+11
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,17 @@ public CompletionStage<R> getReactiveSingleResult() {
139139
return selectionQueryDelegate.getReactiveSingleResult();
140140
}
141141

142+
@Override
143+
public long getResultCount() {
144+
throw LOG.nonReactiveMethodCall( "getReactiveResultCount()" );
145+
}
146+
147+
@Override
148+
public CompletionStage<Long> getReactiveResultCount() {
149+
return selectionQueryDelegate
150+
.getReactiveResultsCount( ( (SqmSelectStatement<?>) getSqmStatement() ).createCountQuery(), this );
151+
}
152+
142153
@Override
143154
public CompletionStage<R> getReactiveSingleResultOrNull() {
144155
return selectionQueryDelegate.getReactiveSingleResultOrNull();

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

+6
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

+2-1
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

+6
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

+46
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+
}

0 commit comments

Comments
 (0)