From efe0785829aef6707f6c82d5f28b13e1df208da6 Mon Sep 17 00:00:00 2001
From: agrancaric <agrancaric@croz.net>
Date: Thu, 22 Aug 2024 09:22:02 +0200
Subject: [PATCH] Add support for sorting by subclass properties when going
 through Spring (Hibernate already supports it but Spring fails with
 exception)

---
 .../nrich/search/support/JpaQueryBuilder.java |  4 +-
 .../jpa/repository/query/NrichQueryUtils.java | 99 +++++++++++++++++++
 .../support/JpaQueryBuilderTest.java          | 13 +++
 3 files changed, 114 insertions(+), 2 deletions(-)
 create mode 100644 nrich-search/src/main/java/org/springframework/data/jpa/repository/query/NrichQueryUtils.java

diff --git a/nrich-search/src/main/java/net/croz/nrich/search/support/JpaQueryBuilder.java b/nrich-search/src/main/java/net/croz/nrich/search/support/JpaQueryBuilder.java
index 00255718..876b862d 100644
--- a/nrich-search/src/main/java/net/croz/nrich/search/support/JpaQueryBuilder.java
+++ b/nrich-search/src/main/java/net/croz/nrich/search/support/JpaQueryBuilder.java
@@ -32,7 +32,7 @@
 import net.croz.nrich.search.util.PathResolvingUtil;
 import net.croz.nrich.search.util.ProjectionListResolverUtil;
 import org.springframework.data.domain.Sort;
-import org.springframework.data.jpa.repository.query.QueryUtils;
+import org.springframework.data.jpa.repository.query.NrichQueryUtils;
 import org.springframework.data.util.DirectFieldAccessFallbackBeanWrapper;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -101,7 +101,7 @@ public <R, P> CriteriaQuery<P> buildQuery(R request, SearchConfiguration<T, P, R
         query.distinct(searchConfiguration.isDistinct());
 
         if (sort != null && sort.isSorted()) {
-            query.orderBy(QueryUtils.toOrders(sort, root, criteriaBuilder));
+            query.orderBy(NrichQueryUtils.toOrders(sort, root, criteriaBuilder));
         }
 
         return query;
diff --git a/nrich-search/src/main/java/org/springframework/data/jpa/repository/query/NrichQueryUtils.java b/nrich-search/src/main/java/org/springframework/data/jpa/repository/query/NrichQueryUtils.java
new file mode 100644
index 00000000..fc3d9687
--- /dev/null
+++ b/nrich-search/src/main/java/org/springframework/data/jpa/repository/query/NrichQueryUtils.java
@@ -0,0 +1,99 @@
+package org.springframework.data.jpa.repository.query;
+
+import org.hibernate.metamodel.model.domain.EntityDomainType;
+import org.springframework.data.domain.Sort;
+import org.springframework.data.mapping.PropertyPath;
+import org.springframework.data.mapping.PropertyReferenceException;
+import org.springframework.util.Assert;
+import org.springframework.util.CollectionUtils;
+import org.springframework.util.ConcurrentReferenceHashMap;
+
+import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.Expression;
+import jakarta.persistence.criteria.From;
+import jakarta.persistence.criteria.Order;
+import jakarta.persistence.metamodel.Type;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Adds support for sorting by subclass properties. Although Hibernate supports sorting by subclass properties, Spring will throw an exception.
+ */
+public final class NrichQueryUtils {
+
+    private static final Map<PropertyPathKey, PropertyPath> PROPERTY_PATH_CACHE = new ConcurrentReferenceHashMap<>();
+
+    private NrichQueryUtils() {
+    }
+
+    public static List<Order> toOrders(Sort sort, From<?, ?> from, CriteriaBuilder cb) {
+        if (sort.isUnsorted()) {
+            return Collections.emptyList();
+        }
+
+        Assert.notNull(from, "From must not be null");
+        Assert.notNull(cb, "CriteriaBuilder must not be null");
+
+        List<jakarta.persistence.criteria.Order> orders = new ArrayList<>();
+
+        for (org.springframework.data.domain.Sort.Order order : sort) {
+            orders.add(toJpaOrder(order, from, cb));
+        }
+
+        return orders;
+    }
+
+    private static jakarta.persistence.criteria.Order toJpaOrder(Sort.Order order, From<?, ?> from, CriteriaBuilder cb) {
+        String propertyName = order.getProperty();
+
+        PropertyPath property = PROPERTY_PATH_CACHE.computeIfAbsent(new PropertyPathKey(propertyName, from.getJavaType()), key -> resolvePropertyPath(propertyName, from));
+        Expression<?> expression = QueryUtils.toExpressionRecursively(from, property);
+
+        if (order.isIgnoreCase() && String.class.equals(expression.getJavaType())) {
+            @SuppressWarnings("unchecked")
+            Expression<String> upper = cb.lower((Expression<String>) expression);
+            return order.isAscending() ? cb.asc(upper) : cb.desc(upper);
+        }
+
+        return order.isAscending() ? cb.asc(expression) : cb.desc(expression);
+    }
+
+    private static PropertyPath resolvePropertyPath(String property, From<?, ?> from) {
+        PropertyReferenceException originalException;
+        try {
+            return PropertyPath.from(property, from.getJavaType());
+        }
+        catch (PropertyReferenceException exception) {
+            originalException = exception;
+        }
+
+        List<? extends Class<?>> subtypeList = resolveSubtypeList(from);
+        if (!CollectionUtils.isEmpty(subtypeList)) {
+            for (Class<?> subtype : subtypeList) {
+                try {
+                    return PropertyPath.from(property, subtype);
+                }
+                catch (PropertyReferenceException exception) {
+                    // ignored
+                }
+            }
+        }
+
+        throw originalException;
+    }
+
+    private static List<? extends Class<?>> resolveSubtypeList(From<?, ?> from) {
+        if (from.getModel() instanceof EntityDomainType<?> entityDomainType) {
+            return entityDomainType.getSubTypes().stream()
+                .map(Type::getJavaType)
+                .toList();
+        }
+
+        return Collections.emptyList();
+    }
+
+    record PropertyPathKey(String property, Class<?> type) {
+    }
+}
diff --git a/nrich-search/src/test/java/net/croz/nrich/search/repository/support/JpaQueryBuilderTest.java b/nrich-search/src/test/java/net/croz/nrich/search/repository/support/JpaQueryBuilderTest.java
index 682a1ef4..4b572f88 100644
--- a/nrich-search/src/test/java/net/croz/nrich/search/repository/support/JpaQueryBuilderTest.java
+++ b/nrich-search/src/test/java/net/croz/nrich/search/repository/support/JpaQueryBuilderTest.java
@@ -389,6 +389,19 @@ void shouldSortEntities() {
         assertThat(results.get(0).getAge()).isEqualTo(28);
     }
 
+    @Test
+    void shouldSortEntitiesBySubclassProperty() {
+        // given
+        TestEntitySearchRequest request = new TestEntitySearchRequest(null);
+
+        // when
+        List<TestEntity> results = executeQuery(request, SearchConfiguration.emptyConfiguration(), Sort.by(Sort.Order.asc("subName")));
+
+        // then
+        assertThat(results).isNotEmpty();
+        assertThat(results.get(0).getNestedEntity().getNestedEntityName()).isEqualTo("nested0");
+    }
+
     @Test
     void shouldSortByJoinedEntity() {
         // given