Skip to content

Commit

Permalink
ORM/HR/Panache: rely on ORM getResultCount
Browse files Browse the repository at this point in the history 
This allows us to not care about which parameters are set on the count
query, since `order by` parameters are ignored by ORM. Which is not
the case if we write our own count query.

Fixes quarkusio#40962
  • Loading branch information
@FroMage
FroMage committed Oct 18, 2024
1 parent 1961df9 commit 9305609
Show file tree
Hide file tree
Showing 12 changed files with 114 additions and 221 deletions.
56 changes: 27 additions & 29 deletions ...src/main/java/io/quarkus/hibernate/orm/panache/common/runtime/CommonPanacheQueryImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,12 @@ public void close() {
* this is the original Panache-Query, if any (can be null)
*/
private String originalQuery;
protected String countQuery;
/**
* This is only used by the Spring Data JPA extension, due to Spring's Query annotation allowing a custom count query
* See https://docs.spring.io/spring-data/jpa/reference/jpa/query-methods.html#jpa.query-methods.at-query.native
* Otherwise we do not use this, and rely on ORM to generate count queries
*/
protected String customCountQueryForSpring;
private String orderBy;
private Session session;

Expand All @@ -73,11 +78,12 @@ public CommonPanacheQueryImpl(Session session, String query, String originalQuer
this.paramsArrayOrMap = paramsArrayOrMap;
}

private CommonPanacheQueryImpl(CommonPanacheQueryImpl<?> previousQuery, String newQueryString, String countQuery,
private CommonPanacheQueryImpl(CommonPanacheQueryImpl<?> previousQuery, String newQueryString,
String customCountQueryForSpring,
Class<?> projectionType) {
this.session = previousQuery.session;
this.query = newQueryString;
this.countQuery = countQuery;
this.customCountQueryForSpring = customCountQueryForSpring;
this.orderBy = previousQuery.orderBy;
this.paramsArrayOrMap = previousQuery.paramsArrayOrMap;
this.page = previousQuery.page;
Expand Down Expand Up @@ -106,16 +112,16 @@ public <T> CommonPanacheQueryImpl<T> project(Class<T> type) {

// If the query starts with a select clause, we pass it on to ORM which can handle that via a projection type
if (lowerCasedTrimmedQuery.startsWith("select ")) {
// just pass it through
return new CommonPanacheQueryImpl<>(this, query, countQuery, type);
// I think projections do not change the result count, so we can keep the custom count query
return new CommonPanacheQueryImpl<>(this, query, customCountQueryForSpring, type);
}

// FIXME: this assumes the query starts with "FROM " probably?

// build select clause with a constructor expression
String selectClause = "SELECT " + getParametersFromClass(type, null);
return new CommonPanacheQueryImpl<>(this, selectClause + selectQuery,
"select count(*) " + selectQuery, null);
// I think projections do not change the result count, so we can keep the custom count query
return new CommonPanacheQueryImpl<>(this, selectClause + selectQuery, customCountQueryForSpring, null);
}

private StringBuilder getParametersFromClass(Class<?> type, String parentParameter) {
Expand Down Expand Up @@ -267,35 +273,27 @@ public void withHint(String hintName, Object value) {

// Results

@SuppressWarnings("unchecked")
public long count() {
if (count == null) {
String selectQuery = query;
if (PanacheJpaUtil.isNamedQuery(query)) {
SelectionQuery q = session.createNamedSelectionQuery(query.substring(1));
selectQuery = getQueryString(q);
}

SelectionQuery countQuery = session.createSelectionQuery(countQuery(selectQuery));
if (paramsArrayOrMap instanceof Map)
AbstractJpaOperations.bindParameters(countQuery, (Map<String, Object>) paramsArrayOrMap);
else
AbstractJpaOperations.bindParameters(countQuery, (Object[]) paramsArrayOrMap);
try (NonThrowingCloseable c = applyFilters()) {
count = (Long) countQuery.getSingleResult();
if (customCountQueryForSpring != null) {
SelectionQuery<Long> countQuery = session.createSelectionQuery(customCountQueryForSpring, Long.class);
if (paramsArrayOrMap instanceof Map)
AbstractJpaOperations.bindParameters(countQuery, (Map<String, Object>) paramsArrayOrMap);
else
AbstractJpaOperations.bindParameters(countQuery, (Object[]) paramsArrayOrMap);
try (NonThrowingCloseable c = applyFilters()) {
count = countQuery.getSingleResult();
}
} else {
SelectionQuery<?> query = createBaseQuery();
try (NonThrowingCloseable c = applyFilters()) {
count = query.getResultCount();
}
}
}
return count;
}

private String countQuery(String selectQuery) {
if (countQuery != null) {
return countQuery;
}

return PanacheJpaUtil.getFastCountQuery(selectQuery);
}

@SuppressWarnings("unchecked")
public <T extends Entity> List<T> list() {
SelectionQuery hibernateQuery = createQuery();
Expand Down
2 changes: 1 addition & 1 deletion ...ntime/src/main/java/io/quarkus/hibernate/orm/panache/runtime/CustomCountPanacheQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ public CustomCountPanacheQuery(Session session, SelectionQuery hibernateQuery, S
super(new CommonPanacheQueryImpl<>(session, CommonPanacheQueryImpl.getQueryString(hibernateQuery),
null, null, paramsArrayOrMap) {
{
this.countQuery = customCountQuery;
this.customCountQueryForSpring = customCountQuery;
}
});
}
Expand Down
52 changes: 25 additions & 27 deletions ...ain/java/io/quarkus/hibernate/reactive/panache/common/runtime/CommonPanacheQueryImpl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,12 @@ public class CommonPanacheQueryImpl<Entity> {
* this is the original Panache-Query, if any (can be null)
*/
private String originalQuery;
protected String countQuery;
/**
* This is only used by the Spring Data JPA extension, due to Spring's Query annotation allowing a custom count query
* See https://docs.spring.io/spring-data/jpa/reference/jpa/query-methods.html#jpa.query-methods.at-query.native
* Otherwise we do not use this, and rely on ORM to generate count queries
*/
protected String customCountQueryForSpring;
private String orderBy;
private Uni<Mutiny.Session> em;

Expand All @@ -62,11 +67,12 @@ public CommonPanacheQueryImpl(Uni<Mutiny.Session> em, String query, String origi
this.paramsArrayOrMap = paramsArrayOrMap;
}

private CommonPanacheQueryImpl(CommonPanacheQueryImpl<?> previousQuery, String newQueryString, String countQuery,
private CommonPanacheQueryImpl(CommonPanacheQueryImpl<?> previousQuery, String newQueryString,
String customCountQueryForSpring,
Class<?> projectionType) {
this.em = previousQuery.em;
this.query = newQueryString;
this.countQuery = countQuery;
this.customCountQueryForSpring = customCountQueryForSpring;
this.orderBy = previousQuery.orderBy;
this.paramsArrayOrMap = previousQuery.paramsArrayOrMap;
this.page = previousQuery.page;
Expand Down Expand Up @@ -94,16 +100,16 @@ public <T> CommonPanacheQueryImpl<T> project(Class<T> type) {

// If the query starts with a select clause, we pass it on to ORM which can handle that via a projection type
if (lowerCasedTrimmedQuery.startsWith("select ")) {
// just pass it through
return new CommonPanacheQueryImpl<>(this, query, countQuery, type);
// I think projections do not change the result count, so we can keep the custom count query
return new CommonPanacheQueryImpl<>(this, query, customCountQueryForSpring, type);
}

// FIXME: this assumes the query starts with "FROM " probably?

// build select clause with a constructor expression
String selectClause = "SELECT " + getParametersFromClass(type, null);
return new CommonPanacheQueryImpl<>(this, selectClause + selectQuery,
"select count(*) " + selectQuery, type);
// I think projections do not change the result count, so we can keep the custom count query
return new CommonPanacheQueryImpl<>(this, selectClause + selectQuery, customCountQueryForSpring, null);
}

private StringBuilder getParametersFromClass(Class<?> type, String parentParameter) {
Expand Down Expand Up @@ -263,34 +269,26 @@ public void withHint(String hintName, Object value) {

@SuppressWarnings("unchecked")
public Uni<Long> count() {
String selectQuery;
if (PanacheJpaUtil.isNamedQuery(query)) {
selectQuery = NamedQueryUtil.getNamedQuery(query.substring(1));
} else {
selectQuery = query;
}

if (count == null) {
// FIXME: question about caching the result here
count = em.flatMap(session -> {
Mutiny.SelectionQuery<Long> countQuery = session.createSelectionQuery(countQuery(selectQuery), Long.class);
if (paramsArrayOrMap instanceof Map)
AbstractJpaOperations.bindParameters(countQuery, (Map<String, Object>) paramsArrayOrMap);
else
AbstractJpaOperations.bindParameters(countQuery, (Object[]) paramsArrayOrMap);
return applyFilters(session, () -> countQuery.getSingleResult());
if (customCountQueryForSpring != null) {
Mutiny.SelectionQuery<Long> countQuery = session.createSelectionQuery(customCountQueryForSpring,
Long.class);
if (paramsArrayOrMap instanceof Map)
AbstractJpaOperations.bindParameters(countQuery, (Map<String, Object>) paramsArrayOrMap);
else
AbstractJpaOperations.bindParameters(countQuery, (Object[]) paramsArrayOrMap);
return applyFilters(session, () -> countQuery.getSingleResult());
} else {
Mutiny.SelectionQuery<?> query = createBaseQuery(session);
return applyFilters(session, () -> query.getResultCount());
}
});
}
return count;
}

private String countQuery(String selectQuery) {
if (countQuery != null) {
return countQuery;
}
return PanacheJpaUtil.getFastCountQuery(selectQuery);
}

@SuppressWarnings({ "unchecked", "rawtypes" })
public <T extends Entity> Uni<List<T>> list() {
return em.flatMap(session -> {
Expand Down
2 changes: 1 addition & 1 deletion .../src/main/java/io/quarkus/hibernate/reactive/panache/runtime/CustomCountPanacheQuery.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ public CustomCountPanacheQuery(Uni<Mutiny.Session> em, String query, String cust
Object paramsArrayOrMap) {
super(new CommonPanacheQueryImpl<Entity>(em, query, null, null, paramsArrayOrMap) {
{
this.countQuery = customCountQuery;
this.customCountQueryForSpring = customCountQuery;
}
});
}
Expand Down
81 changes: 0 additions & 81 deletions ...mon/runtime/src/main/java/io/quarkus/panache/hibernate/common/runtime/PanacheJpaUtil.java
View file
Original file line number Diff line number Diff line change
@@ -1,15 +1,8 @@
package io.quarkus.panache.hibernate.common.runtime;

import java.util.Locale;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.hibernate.grammars.hql.HqlLexer;
import org.hibernate.grammars.hql.HqlParser;
import org.hibernate.grammars.hql.HqlParser.SelectStatementContext;

import io.quarkus.panache.common.Sort;
import io.quarkus.panache.common.exception.PanacheQueryException;

Expand All @@ -36,80 +29,6 @@ public class PanacheJpaUtil {
static final Pattern WITH_PATTERN = Pattern.compile("^\\s*WITH\\s+.*",
Pattern.CASE_INSENSITIVE | Pattern.DOTALL);

/**
* This turns an HQL (already expanded from Panache-QL) query into a count query, using text manipulation
* if we can, because it's faster, or fall back to using the ORM HQL parser in {@link #getCountQueryUsingParser(String)}
*/
public static String getFastCountQuery(String query) {
// try to generate a good count query from the existing query
String countQuery;
// there are no fast ways to get rid of fetches, or WITH
if (FETCH_PATTERN.matcher(query).matches()
|| WITH_PATTERN.matcher(query).matches()) {
return getCountQueryUsingParser(query);
}
// if it starts with select, we can optimise
Matcher selectMatcher = SELECT_PATTERN.matcher(query);
if (selectMatcher.matches()) {
// this one cannot be null
String firstSelection = selectMatcher.group(1).trim().toLowerCase(Locale.ROOT);
if (firstSelection.startsWith("distinct")) {
// if firstSelection matched distinct only, we have something wrong in our selection list, probably functions/parens
// so bail out
if (firstSelection.length() == 8) {
return getCountQueryUsingParser(query);
}
// this one can be null
String secondSelection = selectMatcher.group(2);
// we can only count distinct single columns
if (secondSelection != null && !secondSelection.trim().isEmpty()) {
throw new PanacheQueryException("Count query not supported for select query: " + query);
}
countQuery = "SELECT COUNT(" + firstSelection + ") " + selectMatcher.group(3);
} else {
// it's not distinct, forget the column list
countQuery = "SELECT COUNT(*) " + selectMatcher.group(3);
}
} else if (LONE_SELECT_PATTERN.matcher(query).matches()) {
// a select anywhere else in there might be tricky
return getCountQueryUsingParser(query);
} else if (FROM_PATTERN.matcher(query).matches()) {
countQuery = "SELECT COUNT(*) " + query;
} else {
throw new PanacheQueryException("Count query not supported for select query: " + query);
}

// remove the order by clause
String lcQuery = countQuery.toLowerCase();
int orderByIndex = lcQuery.lastIndexOf(" order by ");
if (orderByIndex != -1) {
countQuery = countQuery.substring(0, orderByIndex);
}
return countQuery;
}

/**
* This turns an HQL (already expanded from Panache-QL) query into a count query, using the
* ORM HQL parser. Slow version, see {@link #getFastCountQuery(String)} for the fast version.
*/
public static String getCountQueryUsingParser(String query) {
HqlLexer lexer = new HqlLexer(CharStreams.fromString(query));
CommonTokenStream tokens = new CommonTokenStream(lexer);
HqlParser parser = new HqlParser(tokens);
SelectStatementContext statement = parser.selectStatement();
try {
CountParserVisitor visitor = new CountParserVisitor();
statement.accept(visitor);
return visitor.result();
} catch (RequiresSubqueryException x) {
// no luck
SubQueryAliasParserVisitor visitor = new SubQueryAliasParserVisitor();
statement.accept(visitor);
String ret = visitor.result();
return "select count( * ) from ( " + ret + " )";
}
}

public static String getEntityName(Class<?> entityClass) {
// FIXME: not true?
// Escape the entity name just in case some keywords are used
Expand Down
82 changes: 0 additions & 82 deletions ...e-common/runtime/src/test/java/io/quarkus/panache/hibernate/common/runtime/CountTest.java
View file

This file was deleted.

Loading

0 comments on commit 9305609

Please sign in to comment.