JPA Criteria使用Treat预加载特定类型关联时遇懒加载错误

我在使用JPA Criteria处理多态关联时遇到问题：orderItems关联包含StandardOrderItem和SpecificOrderItem两种类型，为避免LazyInitialization异常，我希望仅对StandardOrderItem预加载其components关联。但criteriaBuilder.treat()仅支持Join（懒加载），而Fetch（急加载）无法使用treat，导致关联components时出错。

预期实现的调用方式如下：

orderService.findOrderByIdAndEntityGraph(productionOrderId,
  Map.of("orderItems",
    new JPAUtils.JoinNode(
      StandardOrderItem.class,
      Map.of("components", new JPAUtils.JoinNode(EMPTY_MAP))
      )));

现有实现代码 #

public static <T> T getBySpecificationWithGraph(
    EntityManager entityManager,
    Specification<T> specification,
    Map<String, JoinNode> joins,
    Class<T> clazz
  ) {
    CriteriaQuery<T> criteriaQuery = prepareCriteriaQuery(entityManager, specification, joins, clazz);
    return entityManager.createQuery(criteriaQuery).getSingleResult();
  }

  public static <T, V> Specification<T> fieldEquals(String fieldName, V value) {
    return (root, criteriaQuery, criteriaBuilder) ->
      criteriaBuilder.equal(root.get(fieldName), value);
  }

  private static <T> CriteriaQuery<T> prepareCriteriaQuery(
    EntityManager entityManager,
    Specification<T> specification,
    Map<String, JoinNode> joins,
    Class<T> clazz
  ) {
    CriteriaBuilder criteriaBuilder = entityManager.getCriteriaBuilder();
    CriteriaQuery<T> criteriaQuery = criteriaBuilder.createQuery(clazz);
    Root<T> ownerRoot = criteriaQuery.from(clazz);

    // Apply fetch logic with TREAT
    prepareFetch(joins, ownerRoot, null, criteriaBuilder);

    criteriaQuery.select(ownerRoot)
      .where(specification.toPredicate(ownerRoot, criteriaQuery, criteriaBuilder));
    return criteriaQuery;
  }

  private static <T> void prepareFetch(
    Map<String, JoinNode> joins,
    From<?, T> root,  // Use From to accommodate both Root and Join
    Fetch<?, ?> fetch,
    CriteriaBuilder criteriaBuilder
  ) {
    joins.forEach((key, value) -> {
      Fetch<?, ?> childFetch = null;
      Join<?, ?> childJoin = null;  // Declare childJoin for TREAT usage

      // Apply TREAT logic only on Fetches or Joins
      if (fetch == null) {
        if (value.getType() != null) {
          // First, attempt to treat it as a collection join
          try {
            // If the relationship is a collection (e.g., @OneToMany)
            SetJoin<Object, Object> setJoin = root.joinSet(key, JoinType.LEFT);
            childJoin = criteriaBuilder.treat(setJoin, value.getType());

            root.fetch(key, JoinType.LEFT);  // Ensure the association is fetched eagerly
          } catch (IllegalArgumentException | ClassCastException e) {
            // If it's not a collection, treat it as a singular join (e.g., @ManyToOne, @OneToOne)
            try {
              Join<Object, Object> join = root.join(key, JoinType.LEFT);
              childJoin = criteriaBuilder.treat(join, value.getType());

              root.fetch(key, JoinType.LEFT);  // Ensure the association is fetched eagerly
            } catch (IllegalArgumentException | ClassCastException ex) {
              throw new RuntimeException("Unable to apply TREAT to the join: " + key, ex);
            }
          }
        } else {
          // If no type is provided, just fetch the relationship
          childFetch = root.fetch(key, JoinType.LEFT);
        }
      } else {
        childFetch = fetch.fetch(key, JoinType.LEFT);
      }

      // Handle recursive joins and fetches for deeper relationships
      Objects.requireNonNull(value);
      Map<String, JoinNode> next = value.getJoins();

      if (!next.isEmpty()) {
        // Use childJoin for TREAT case or childFetch for normal cases
        if (childJoin != null) {
          prepareFetch(next, childJoin, null, criteriaBuilder);  // Continue joining with TREAT
        } else {
          prepareFetch(next, root, childFetch, criteriaBuilder);  // Continue fetching
        }
      }
    });
  }

  @Getter
  public static class JoinNode {
    private final Class<?> type;
    private final Map<String, JoinNode> joins;

    public JoinNode(Class<?> type, Map<String, JoinNode> joins) {
      this.type = type;
      this.joins = joins;
    }

    public JoinNode(Map<String, JoinNode> joins) {
      this.joins = joins;
      this.type = null;
    }
  }

使用示例 #

orderService.findOrderByIdAndEntityGraph(productionOrderId,
  Map.of("orderItems",
    new JPAUtils.JoinNode(
      StandardOrderItem.class,
      Map.of("components", new JPAUtils.JoinNode(EMPTY_MAP))
      )));

解决方案 #

现有代码的核心问题是同时调用root.join()和root.fetch()导致重复关联，且treat后的Join未与Fetch关联，无法实现子类关联的预加载。以下是三种可行解决方式：

方式一：基于Hibernate扩展API实现（适配Hibernate作为JPA实现） #

Hibernate的FetchImplementor接口可将Fetch转换为Join，从而支持treat操作，修改prepareFetch方法如下：

private static <T> void prepareFetch(
    Map<String, JoinNode> joins,
    From<?, T> root,
    Fetch<?, ?> fetch,
    CriteriaBuilder criteriaBuilder
) {
    joins.forEach((key, value) -> {
        Fetch<?, ?> childFetch = null;
        Join<?, ?> childJoin = null;

        if (fetch == null) {
            if (value.getType() != null) {
                // 获取Fetch并转换为Join（Hibernate扩展）
                Fetch<?, ?> tempFetch = root.fetch(key, JoinType.LEFT);
                Join<?, ?> tempJoin = ((FetchImplementor<?, ?>) tempFetch).toJoin();
                childJoin = criteriaBuilder.treat(tempJoin, value.getType());
            } else {
                childFetch = root.fetch(key, JoinType.LEFT);
            }
        } else {
            if (value.getType() != null) {
                Fetch<?, ?> tempFetch = fetch.fetch(key, JoinType.LEFT);
                Join<?, ?> tempJoin = ((FetchImplementor<?, ?>) tempFetch).toJoin();
                childJoin = criteriaBuilder.treat(tempJoin, value.getType());
            } else {
                childFetch = fetch.fetch(key, JoinType.LEFT);
            }
        }

        Map<String, JoinNode> next = value.getJoins();
        if (!next.isEmpty()) {
            if (childJoin != null) {
                next.forEach((subKey, subValue) -> {
                    if (subValue.getType() != null) {
                        Join<?, ?> subJoin = criteriaBuilder.treat(childJoin.join(subKey, JoinType.LEFT), subValue.getType());
                        childJoin.fetch(subKey, JoinType.LEFT);
                        prepareFetch(subValue.getJoins(), subJoin, null, criteriaBuilder);
                    } else {
                        Fetch<?, ?> subFetch = childJoin.fetch(subKey, JoinType.LEFT);
                        prepareFetch(subValue.getJoins(), childJoin, subFetch, criteriaBuilder);
                    }
                });
            } else {
                prepareFetch(next, root, childFetch, criteriaBuilder);
            }
        }
    });
}

方式二：使用JPA标准实体图（推荐） #

通过定义带子类过滤的实体图，无需手动处理Criteria逻辑，符合JPA标准：

// 在Order实体上定义命名实体图
@NamedEntityGraph(name = "Order.withStandardOrderItemsComponents",
    attributeNodes = {
        @NamedAttributeNode(value = "orderItems", subgraph = "standardOrderItem")
    },
    subgraphs = {
        @NamedSubgraph(name = "standardOrderItem",
            type = StandardOrderItem.class,
            attributeNodes = @NamedAttributeNode("components"))
    }
)
@Entity
public class Order {
    // 实体字段...
    @OneToMany(mappedBy = "order")
    private List<OrderItem> orderItems;
}

// 查询时使用实体图
public Order findOrderByIdAndEntityGraph(Long id) {
    EntityGraph<Order> entityGraph = entityManager.getEntityGraph("Order.withStandardOrderItemsComponents");
    return entityManager.find(Order.class, id, Collections.singletonMap("javax.persistence.fetchgraph", entityGraph));
}

方式三：修正现有Criteria逻辑 #

确保treat后的Join与Fetch关联，避免重复操作：

private static <T> void prepareFetch(
    Map<String, JoinNode> joins,
    From<?, T> root,
    Fetch<?, ?> fetch,
    CriteriaBuilder criteriaBuilder
) {
    joins.forEach((key, value) -> {
        From<?, ?> currentFrom = null;
        Fetch<?, ?> currentFetch = null;

        if (fetch == null) {
            if (value.getType() != null) {
                // 先Join并treat，再基于该Join执行Fetch
                Join<?, ?> join = root.join(key, JoinType.LEFT);
                currentFrom = criteriaBuilder.treat(join, value.getType());
                root.fetch(key, JoinType.LEFT);
            } else {
                currentFetch = root.fetch(key, JoinType.LEFT);
                currentFrom = root;
            }
        } else {
            if (value.getType() != null) {
                Join<?, ?> join = ((FetchImplementor<?, ?>) fetch).toJoin();
                currentFrom = criteriaBuilder.treat(join.join(key, JoinType.LEFT), value.getType());
                fetch.fetch(key, JoinType.LEFT);
            } else {
                currentFetch = fetch.fetch(key, JoinType.LEFT);
                currentFrom = ((FetchImplementor<?, ?>) fetch).toJoin();
            }
        }

        Map<String, JoinNode> next = value.getJoins();
        if (!next.isEmpty()) {
            if (currentFrom instanceof Join) {
                prepareFetch(next, (Join<?, ?>) currentFrom, null, criteriaBuilder);
            } else {
                prepareFetch(next, (From<?, T>) currentFrom, currentFetch, criteriaBuilder);
            }
        }
    });
}

内容的提问来源于stack exchange，提问作者Yurii Ripetskyi