Hibernate 7: give criterion handlers the query context and let a handler order by an expression (8.1.x)

grails-data-hibernate7/core/src/main/groovy/org/grails/orm/hibernate/query/JpaCriteriaQueryCreator.java

@@ -274,30 +274,41 @@ private Class<?> resolveProjectionType(Query.PropertyProjection projection) {
     private void assignOrderBy(CriteriaQuery<?> cq, JpaQueryContext context) {
         if (isCountOnly()) {
             // a count returns one row whatever the order, and an ORDER BY on a column that is not aggregated fails
-            // on H2 and most other databases, so a count query leaves the criteria's orders out
+            // on H2 and most other databases, so a count query leaves the criteria's orders out, handled ones too
             return;
         }
         List<Query.Order> orders = detachedCriteria.getOrders();
         if (!orders.isEmpty()) {
-            var jpaOrders = orders.stream()
-                    .map(order -> {
-                        var propertyName = order.getProperty();
-                        SqlProjection sqlProjection = findSqlProjection(propertyName, false);
-                        Expression<?> expression = sqlProjection != null ?
-                                context.getSelectionAlias(sqlProjection.getColumnAlias()) :
-                                context.getFullyQualifiedExpression(propertyName);
-                        if (order.isIgnoreCase() && expression.getJavaType().equals(String.class)) {
-                            return order.getDirection().equals(Query.Order.Direction.ASC) ?
-                                    criteriaBuilder.asc(criteriaBuilder.lower((Expression<String>) expression)) :
-                                    criteriaBuilder.desc(criteriaBuilder.lower((Expression<String>) expression));
-                        } else {
-                            return order.getDirection().equals(Query.Order.Direction.ASC) ?
-                                    criteriaBuilder.asc(expression) :
-                                    criteriaBuilder.desc(expression);
-                        }
-                    })
-                    .toArray(jakarta.persistence.criteria.Order[]::new);
-            cq.orderBy(jpaOrders);
+            List<jakarta.persistence.criteria.Order> jpaOrders = new ArrayList<>();
+            for (Query.Order order : orders) {
+                PredicateGenerator.OrderHandler customHandler = PredicateGenerator.orderHandler(order.getClass());
+                if (customHandler != null) {
+                    // the handler of a Query.Order subclass builds its JPA order, or leaves the order out
+                    jakarta.persistence.criteria.Order customOrder =
+                            customHandler.handle(cq, context.getRoot(), criteriaBuilder, order, context);
+                    if (customOrder != null) {
+                        jpaOrders.add(customOrder);
+                    }
+                    continue;
+                }
+                String propertyName = order.getProperty();
+                SqlProjection sqlProjection = findSqlProjection(propertyName, false);
+                Expression<?> expression = sqlProjection != null ?
+                        context.getSelectionAlias(sqlProjection.getColumnAlias()) :
+                        context.getFullyQualifiedExpression(propertyName);
+                jakarta.persistence.criteria.Order jpaOrder;
+                if (order.isIgnoreCase() && expression.getJavaType().equals(String.class)) {
+                    jpaOrder = order.getDirection().equals(Query.Order.Direction.ASC) ?
+                            criteriaBuilder.asc(criteriaBuilder.lower((Expression<String>) expression)) :
+                            criteriaBuilder.desc(criteriaBuilder.lower((Expression<String>) expression));
+                } else {
+                    jpaOrder = order.getDirection().equals(Query.Order.Direction.ASC) ?
+                            criteriaBuilder.asc(expression) :
+                            criteriaBuilder.desc(expression);
+                }
+                jpaOrders.add(jpaOrder);
+            }
+            cq.orderBy(jpaOrders.toArray(new jakarta.persistence.criteria.Order[0]));
         }
     }
 

grails-data-hibernate7/core/src/main/groovy/org/grails/orm/hibernate/query/JpaQueryContext.java

@@ -34,7 +34,7 @@
  *
  * @since 8.0
  */
-public class JpaQueryContext implements Cloneable {
+public class JpaQueryContext implements Cloneable, QueryPathContext {
 
     private final AliasRegistry aliasRegistry;
     private final JoinTracker joinTracker;
@@ -110,6 +110,7 @@ public void setParent(JpaQueryContext parent) {
         this.parent = parent;
     }
 
+    @Override
     public From<?, ?> getRoot() {
         return joinTracker.getRoot();
     }
@@ -143,11 +144,13 @@ public void registerSelectionAlias(String alias, Expression<?> selection) {
     }
 
     /**
-     * Returns the selection recorded under an alias by {@link #registerSelectionAlias(String, Expression)}.
+     * Returns the selection recorded under an alias by {@link #registerSelectionAlias(String, Expression)}. Unlike
+     * {@link #getFullyQualifiedExpression(String)}, which resolves a property first, this asks for the alias alone.
      *
      * @param alias the alias of the selection
      * @return the selection, or {@code null} if no selection of this query has the alias
      */
+    @Override
     public Expression<?> getSelectionAlias(String alias) {
         return selectionAliases.get(alias);
     }
@@ -173,6 +176,7 @@ public Expression<?> getAliasedExpression(String alias) {
         return expr;
     }
 
+    @Override
     public Expression<?> getFullyQualifiedExpression(String path) {
         if (parent != null) {
             if ("{alias}".equals(path)) {
@@ -185,6 +189,7 @@ public Expression<?> getFullyQualifiedExpression(String path) {
         return resolver.resolve(path);
     }
 
+    @Override
     public Path<?> getFullyQualifiedPath(String path) {
         if (parent != null) {
             if ("{alias}".equals(path)) {

grails-data-hibernate7/core/src/main/groovy/org/grails/orm/hibernate/query/PredicateGenerator.java

@@ -27,11 +27,13 @@
 import java.util.Objects;
 import java.util.ServiceLoader;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
+import groovy.lang.Closure;
+
 import jakarta.persistence.criteria.AbstractQuery;
 import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
 import jakarta.persistence.criteria.Expression;
 import jakarta.persistence.criteria.From;
 import jakarta.persistence.criteria.JoinType;
@@ -90,6 +92,123 @@ Predicate handle(
             HibernateCriteriaBuilder criteriaBuilder,
             Query.Criterion criterion
         );
+
+        /**
+         * Converts a criterion with the context of the query it belongs to. The context resolves a property path
+         * the way the built-in criteria do: {@link QueryPathContext#getFullyQualifiedPath(String)} turns
+         * {@code alias.property}, for an alias created with {@code createAlias}, or {@code association.property}
+         * into a path on the join GORM created for that association, so a custom criterion and a GORM criterion on
+         * the same association restrict the same row instead of each joining the association again. This default
+         * ignores the context and calls {@link #handle(AbstractQuery, From, HibernateCriteriaBuilder, Query.Criterion)};
+         * a handler that needs the context implements {@link ContextualCriterionHandler}.
+         *
+         * @param criteriaQuery the query or subquery the predicate is for
+         * @param root the root of that query
+         * @param criteriaBuilder the criteria builder
+         * @param criterion the criterion to convert
+         * @param context the context of the query: its root, its joins and its aliases
+         * @return the predicate, or {@code null} to add none
+         * @since 8.1
+         */
+        default Predicate handle(
+            AbstractQuery<?> criteriaQuery,
+            From<?, ?> root,
+            HibernateCriteriaBuilder criteriaBuilder,
+            Query.Criterion criterion,
+            QueryPathContext context
+        ) {
+            return handle(criteriaQuery, root, criteriaBuilder, criterion);
+        }
+    }
+
+    /**
+     * A {@link CriterionHandler} that is given the context of the query, through which it resolves property paths
+     * with {@link QueryPathContext#getFullyQualifiedPath(String)}. A closure with five parameters registered with
+     * {@link #registerCriterionHandler(Class, Closure)} becomes one, with no {@code as}:
+     * <pre>
+     *     PredicateGenerator.registerCriterionHandler(MyCriterion) { query, root, cb, criterion, context -&gt;
+     *         def c = criterion as MyCriterion
+     *         cb.like(context.getFullyQualifiedPath(c.property) as Expression&lt;String&gt;, c.value as String)
+     *     }
+     * </pre>
+     *
+     * @since 8.1
+     */
+    @FunctionalInterface
+    public interface ContextualCriterionHandler extends CriterionHandler {
+
+        @Override
+        Predicate handle(
+            AbstractQuery<?> criteriaQuery,
+            From<?, ?> root,
+            HibernateCriteriaBuilder criteriaBuilder,
+            Query.Criterion criterion,
+            QueryPathContext context
+        );
+
+        /**
+         * Called by code that has no context to give; the handler then receives {@code null} for it.
+         */
+        @Override
+        default Predicate handle(
+            AbstractQuery<?> criteriaQuery,
+            From<?, ?> root,
+            HibernateCriteriaBuilder criteriaBuilder,
+            Query.Criterion criterion
+        ) {
+            return handle(criteriaQuery, root, criteriaBuilder, criterion, null);
+        }
+    }
+
+    /**
+     * Extension point for ordering a criteria query by something other than a property. Register a handler for a
+     * {@link Query.Order} subclass to have it converted to a JPA {@link jakarta.persistence.criteria.Order}, in
+     * the place the order has among the query's orders. The order reaches the query through
+     * {@code order(Query.Order)} of the criteria builder or of a detached criteria.
+     *
+     * <pre>
+     *     class ByRandom extends Query.Order {
+     *         ByRandom() { super('random') }
+     *     }
+     *     PredicateGenerator.registerOrderHandler(ByRandom) { query, root, cb, order, context -&gt;
+     *         cb.asc(cb.function('random', Double))
+     *     }
+     * </pre>
+     *
+     * @since 8.1
+     */
+    @FunctionalInterface
+    public interface OrderHandler {
+
+        /**
+         * @param criteriaQuery the query, whose result type tells a query for the entity from a projection
+         * @param root the root of the query
+         * @param criteriaBuilder the criteria builder
+         * @param order the order to convert
+         * @param context the context of the query: its root, its joins and its aliases
+         * @return the JPA order, or {@code null} to leave this order out of the query
+         */
+        jakarta.persistence.criteria.Order handle(
+            CriteriaQuery<?> criteriaQuery,
+            From<?, ?> root,
+            HibernateCriteriaBuilder criteriaBuilder,
+            Query.Order order,
+            QueryPathContext context
+        );
+    }
+
+    /**
+     * SPI for contributing an {@link OrderHandler} via {@link ServiceLoader}. Register an implementation in
+     * {@code META-INF/services/org.grails.orm.hibernate.query.PredicateGenerator$OrderHandlerProvider}; it is
+     * discovered on first query execution, like a {@link CriterionHandlerProvider}.
+     *
+     * @since 8.1
+     */
+    public interface OrderHandlerProvider {
+
+        Class<? extends Query.Order> orderType();
+
+        OrderHandler orderHandler();
     }
 
     /**
@@ -119,13 +238,69 @@ public interface CriterionHandlerProvider {
     public static final Map<Class<? extends Query.Criterion>, CriterionHandler> CUSTOM_HANDLERS =
         new ConcurrentHashMap<>();
 
-    private static final AtomicBoolean SERVICE_LOADERS_INITIALIZED = new AtomicBoolean(false);
+    private static final Map<Class<? extends Query.Order>, OrderHandler> CUSTOM_ORDER_HANDLERS =
+        new ConcurrentHashMap<>();
+
+    private static final Object SERVICE_LOADER_LOCK = new Object();
+
+    private static volatile boolean serviceLoadersInitialized = false;
 
+    /**
+     * Loads the handler providers once. A caller that arrives while another thread is loading them waits until the
+     * loading is complete, so no query decides on a handler before the discovered ones are registered. Discovery uses
+     * {@code putIfAbsent}, so a discovered handler does not replace one that is already registered; registering in
+     * code uses {@code put}, so a handler registered in code replaces a discovered one.
+     */
     private static void loadServiceProviders() {
-        if (SERVICE_LOADERS_INITIALIZED.compareAndSet(false, true)) {
-            ServiceLoader.load(CriterionHandlerProvider.class,
-                    Thread.currentThread().getContextClassLoader())
+        if (serviceLoadersInitialized) {
+            return;
+        }
+        synchronized (SERVICE_LOADER_LOCK) {
+            if (serviceLoadersInitialized) {
+                return;
+            }
+            ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
+            ServiceLoader.load(CriterionHandlerProvider.class, classLoader)
                 .forEach(p -> CUSTOM_HANDLERS.putIfAbsent(p.criterionType(), p.criterionHandler()));
+            ServiceLoader.load(OrderHandlerProvider.class, classLoader)
+                .forEach(p -> CUSTOM_ORDER_HANDLERS.putIfAbsent(p.orderType(), p.orderHandler()));
+            serviceLoadersInitialized = true;
+        }
+    }
+
+    /**
+     * Registers an {@link OrderHandler} for the given order type. The handler converts every order of exactly
+     * that type; programmatic registration takes precedence over a {@link ServiceLoader}-discovered handler.
+     *
+     * @param type the exact order class to handle
+     * @param handler the handler that converts the order to a JPA order
+     * @since 8.1
+     */
+    public static void registerOrderHandler(Class<? extends Query.Order> type, OrderHandler handler) {
+        CUSTOM_ORDER_HANDLERS.put(type, handler);
+    }
+
+    /**
+     * Returns the handler registered for an order type, loading the {@link ServiceLoader} providers first.
+     *
+     * @param type the exact order class
+     * @return the handler, or {@code null} when the type is ordered by its property as usual
+     * @since 8.1
+     */
+    public static OrderHandler orderHandler(Class<? extends Query.Order> type) {
+        loadServiceProviders();
+        return CUSTOM_ORDER_HANDLERS.get(type);
+    }
+
+    /**
+     * Removes all registered custom order handlers and resets the ServiceLoader flag. Useful for test cleanup.
+     *
+     * @since 8.1
+     */
+    public static void clearCustomOrderHandlers() {
+        synchronized (SERVICE_LOADER_LOCK) {
+            CUSTOM_ORDER_HANDLERS.clear();
+            serviceLoadersInitialized = false;
         }
     }
 
@@ -141,13 +316,36 @@ public static void registerCriterionHandler(Class<? extends Query.Criterion> typ
         CUSTOM_HANDLERS.put(type, handler);
     }
 
+    /**
+     * Registers a closure as the {@link CriterionHandler} for the given criterion type. A closure with five
+     * parameters, {@code query, root, criteriaBuilder, criterion, context}, becomes a
+     * {@link ContextualCriterionHandler}; any other closure becomes a {@link CriterionHandler} and is called with
+     * four arguments. This is what makes the two shapes work without {@code as}, as {@link #registerOrderHandler} does
+     * for an order handler.
+     *
+     * @param type    the exact criterion class to handle
+     * @param handler the closure that converts the criterion to a JPA {@link Predicate}
+     * @since 8.1
+     */
+    public static void registerCriterionHandler(Class<? extends Query.Criterion> type, Closure<?> handler) {
+        if (handler.getMaximumNumberOfParameters() == 5) {
+            registerCriterionHandler(type, (ContextualCriterionHandler) (query, root, criteriaBuilder, criterion, context) ->
+                (Predicate) handler.call(query, root, criteriaBuilder, criterion, context));
+        } else {
+            registerCriterionHandler(type, (CriterionHandler) (query, root, criteriaBuilder, criterion) ->
+                (Predicate) handler.call(query, root, criteriaBuilder, criterion));
+        }
+    }
+
     /**
      * Removes all registered custom criterion handlers and resets the ServiceLoader flag.
      * Useful for test cleanup.
      */
     public static void clearCustomCriterionHandlers() {
-        CUSTOM_HANDLERS.clear();
-        SERVICE_LOADERS_INITIALIZED.set(false);
+        synchronized (SERVICE_LOADER_LOCK) {
+            CUSTOM_HANDLERS.clear();
+            serviceLoadersInitialized = false;
+        }
     }
 
     private final HibernateCriteriaBuilder criteriaBuilder;
@@ -190,7 +388,7 @@ public Predicate handleCriterion(
         if (criterion instanceof Query.Criterion c) {
             CriterionHandler customHandler = CUSTOM_HANDLERS.get(c.getClass());
             if (customHandler != null) {
-                return customHandler.handle(criteriaQuery, root, criteriaBuilder, c);
+                return customHandler.handle(criteriaQuery, root, criteriaBuilder, c, fromsByProvider);
             }
         }