Add JpaSpecificationEntityGraphExecutor (#2120)

Provides option to find entities by Specificaton with EntityGraph filter in order to load some LAZY attributes

Signed-off-by: Avgustin Marinov <[email protected]>

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/management/JpaActionManagement.java

@@ -24,6 +24,7 @@
 import org.eclipse.hawkbit.repository.jpa.builder.JpaActionStatusCreate;
 import org.eclipse.hawkbit.repository.jpa.model.JpaAction;
 import org.eclipse.hawkbit.repository.jpa.model.JpaActionStatus;
+import org.eclipse.hawkbit.repository.jpa.model.JpaAction_;
 import org.eclipse.hawkbit.repository.jpa.repository.ActionRepository;
 import org.eclipse.hawkbit.repository.jpa.repository.ActionStatusRepository;
 import org.eclipse.hawkbit.repository.jpa.specifications.ActionSpecifications;
@@ -32,7 +33,7 @@
 import org.eclipse.hawkbit.repository.model.ActionStatus;
 import org.eclipse.hawkbit.repository.model.TargetUpdateStatus;
 import org.springframework.data.domain.PageRequest;
-import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Sort;
 
 /**
  * Implements utility methods for managing {@link Action}s
@@ -103,17 +104,20 @@ protected void assertActionStatusMessageQuota(final JpaActionStatus actionStatus
     }
 
     protected List<Action> findActiveActionsWithHighestWeightConsideringDefault(final String controllerId, final int maxActionCount) {
-        final Pageable pageable = PageRequest.of(0, maxActionCount);
         return Stream.concat(
                         // get the highest actions with weight
-                        actionRepository
-                                .findFetchDsByTarget_ControllerIdAndActiveIsTrueAndWeightNotNullOrderByWeightDescIdAsc(controllerId, pageable)
-                                .stream(),
+                        actionRepository.findAll(
+                                ActionSpecifications.byTargetControllerIdAndActiveAndWeightIsNull(controllerId, false),
+                                JpaAction_.GRAPH_ACTION_DS,
+                                PageRequest.of(0, maxActionCount, Sort.by(Sort.Order.desc(JpaAction_.WEIGHT), Sort.Order.asc(JpaAction_.ID)))).stream(),
                         // get the oldest actions without weight
-                        actionRepository.findFetchDsByTarget_ControllerIdAndActiveIsTrueAndWeightIsNullOrderByIdAsc(controllerId, pageable)
-                                .stream())
+                        actionRepository.findAll(
+                                ActionSpecifications.byTargetControllerIdAndActiveAndWeightIsNull(controllerId, true),
+                                JpaAction_.GRAPH_ACTION_DS,
+                                PageRequest.of(0, maxActionCount, Sort.by(Sort.Order.asc(JpaAction_.ID)))).stream())
                 .sorted(Comparator.comparingInt(this::getWeightConsideringDefault).reversed().thenComparing(Action::getId))
                 .limit(maxActionCount)
+                .map(Action.class::cast)
                 .toList();
     }
 

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/management/JpaControllerManagement.java

@@ -287,12 +287,15 @@ public Action addUpdateActionStatus(final ActionStatusCreate statusCreate) {
     public Optional<Action> findActiveActionWithHighestWeight(final String controllerId) {
         return Stream.concat(
                         // get the highest action with weight
-                        actionRepository
-                                .findByTarget_ControllerIdAndActiveIsTrueAndWeightNotNullOrderByWeightDescIdAsc(controllerId, PAGEABLE_1)
-                                .stream(),
+                        actionRepository.findAll(
+                                ActionSpecifications.byTargetControllerIdAndActiveAndWeightIsNull(controllerId, false),
+                                PageRequest.of(0, 1, Sort.by(Sort.Order.desc(JpaAction_.WEIGHT), Sort.Order.asc(JpaAction_.ID)))).stream(),
                         // get the oldest action without weight
-                        actionRepository.findByTarget_ControllerIdAndActiveIsTrueAndWeightIsNullOrderByIdAsc(controllerId, PAGEABLE_1).stream())
-                .min(Comparator.comparingInt(this::getWeightConsideringDefault).reversed().thenComparing(Action::getId));
+                        actionRepository.findAll(
+                                ActionSpecifications.byTargetControllerIdAndActiveAndWeightIsNull(controllerId, false),
+                                PageRequest.of(0, 1, Sort.by(Sort.Order.asc(JpaAction_.ID)))).stream())
+                .min(Comparator.comparingInt(this::getWeightConsideringDefault).reversed().thenComparing(Action::getId))
+                .map(Action.class::cast);
     }
 
     @Override

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/repository/ActionRepository.java

@@ -57,14 +57,6 @@ public interface ActionRepository extends BaseEntityRepository<JpaAction> {
     Optional<Action> findFirstByTargetIdAndDistributionSetIdAndStatusOrderByIdDesc(
             @Param("target") long targetId, @Param("ds") Long dsId, @Param("status") Action.Status status);
 
-    List<Action> findByTarget_ControllerIdAndActiveIsTrueAndWeightIsNullOrderByIdAsc(String controllerId, Pageable pageable);
-    List<Action> findByTarget_ControllerIdAndActiveIsTrueAndWeightNotNullOrderByWeightDescIdAsc(String controllerId, Pageable pageable);
-
-    @EntityGraph(value = "Action.ds", type = EntityGraphType.LOAD)
-    List<Action> findFetchDsByTarget_ControllerIdAndActiveIsTrueAndWeightIsNullOrderByIdAsc(String controllerId, Pageable pageable);
-    @EntityGraph(value = "Action.ds", type = EntityGraphType.LOAD)
-    List<Action> findFetchDsByTarget_ControllerIdAndActiveIsTrueAndWeightNotNullOrderByWeightDescIdAsc(String controllerId, Pageable pageable);
-
     /**
      * Switches the status of actions from one specific status into another, only if the actions are in a specific status. This should be
      * an atomic operation.

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/repository/BaseEntityRepository.java

@@ -10,16 +10,20 @@
 package org.eclipse.hawkbit.repository.jpa.repository;
 
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
+import jakarta.persistence.EntityGraph;
 import jakarta.persistence.EntityManager;
 
 import org.eclipse.hawkbit.repository.jpa.acm.AccessController;
 import org.eclipse.hawkbit.repository.jpa.model.AbstractJpaBaseEntity_;
 import org.eclipse.hawkbit.repository.jpa.model.AbstractJpaTenantAwareBaseEntity;
 import org.eclipse.hawkbit.repository.model.TenantAwareBaseEntity;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.domain.Specification;
 import org.springframework.data.jpa.repository.JpaSpecificationExecutor;
 import org.springframework.data.repository.CrudRepository;
@@ -36,8 +40,8 @@
 @NoRepositoryBean
 @Transactional(readOnly = true)
 public interface BaseEntityRepository<T extends AbstractJpaTenantAwareBaseEntity>
-        extends PagingAndSortingRepository<T, Long>, CrudRepository<T, Long>, JpaSpecificationExecutor<T>, NoCountSliceRepository<T>,
-        ACMRepository<T> {
+        extends PagingAndSortingRepository<T, Long>, CrudRepository<T, Long>, JpaSpecificationExecutor<T>,
+        JpaSpecificationEntityGraphExecutor<T>, NoCountSliceRepository<T>, ACMRepository<T> {
 
     /**
      * Overrides {@link org.springframework.data.repository.CrudRepository#saveAll(Iterable)} to return a list of created entities instead

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/repository/BaseEntityRepositoryACM.java

@@ -19,6 +19,7 @@
 import java.util.Set;
 import java.util.function.Function;
 
+import jakarta.persistence.EntityManager;
 import jakarta.transaction.Transactional;
 
 import lombok.extern.slf4j.Slf4j;
@@ -287,6 +288,30 @@ public Class<T> getDomainClass() {
         return repository.getDomainClass();
     }
 
+    @Override
+    public Optional<T> findOne(final Specification<T> spec, final String entityGraph) {
+        return repository.findOne(
+                accessController.appendAccessRules(AccessController.Operation.READ, spec), entityGraph);
+    }
+
+    @Override
+    public List<T> findAll(final Specification<T> spec, final String entityGraph) {
+        return repository.findAll(
+                accessController.appendAccessRules(AccessController.Operation.READ, spec), entityGraph);
+    }
+
+    @Override
+    public Page<T> findAll(final Specification<T> spec, final String entityGraph, final Pageable pageable) {
+        return repository.findAll(
+                accessController.appendAccessRules(AccessController.Operation.READ, spec), entityGraph, pageable);
+    }
+
+    @Override
+    public List<T> findAll(final Specification<T> spec, final String entityGraph, final Sort sort) {
+        return repository.findAll(
+                accessController.appendAccessRules(AccessController.Operation.READ, spec), entityGraph, sort);
+    }
+
     @SuppressWarnings("unchecked")
     static <T extends AbstractJpaTenantAwareBaseEntity, R extends BaseEntityRepository<T>> R of(
             final R repository, @NonNull final AccessController<T> accessController) {

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/repository/HawkbitBaseRepository.java

@@ -10,10 +10,15 @@
 package org.eclipse.hawkbit.repository.jpa.repository;
 
 import java.io.Serializable;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
+import jakarta.persistence.EntityGraph;
 import jakarta.persistence.EntityManager;
+import jakarta.persistence.NoResultException;
+import jakarta.persistence.Query;
 import jakarta.persistence.TypedQuery;
 import jakarta.transaction.Transactional;
 
@@ -23,11 +28,13 @@
 import org.springframework.data.domain.PageImpl;
 import org.springframework.data.domain.Pageable;
 import org.springframework.data.domain.Slice;
+import org.springframework.data.domain.Sort;
 import org.springframework.data.jpa.domain.Specification;
 import org.springframework.data.jpa.repository.support.JpaEntityInformation;
 import org.springframework.data.jpa.repository.support.SimpleJpaRepository;
 import org.springframework.lang.NonNull;
 import org.springframework.lang.Nullable;
+import org.springframework.util.ObjectUtils;
 
 /**
  * Repository implementation that allows findAll with disabled count query.
@@ -36,14 +43,18 @@
  * @param <ID> the type of the id of the entity the repository manages
  */
 public class HawkbitBaseRepository<T, ID extends Serializable> extends SimpleJpaRepository<T, ID>
-        implements NoCountSliceRepository<T>, ACMRepository<T> {
+        implements JpaSpecificationEntityGraphExecutor<T>, NoCountSliceRepository<T>, ACMRepository<T> {
 
-    public HawkbitBaseRepository(final Class<T> domainClass, final EntityManager em) {
-        super(domainClass, em);
+    private final EntityManager entityManager;
+
+    public HawkbitBaseRepository(final Class<T> domainClass, final EntityManager entityManager) {
+        super(domainClass, entityManager);
+        this.entityManager = entityManager;
     }
 
     public HawkbitBaseRepository(final JpaEntityInformation<T, ?> entityInformation, final EntityManager entityManager) {
         super(entityInformation, entityManager);
+        this.entityManager = entityManager;
     }
 
     @Override
@@ -93,6 +104,32 @@ public long count(@Nullable final AccessController.Operation operation, @Nullabl
         return count(spec);
     }
 
+    @Override
+    public Optional<T> findOne(final Specification<T> spec, final String entityGraph) {
+        try {
+            return Optional.of(withEntityGraph(getQuery(spec, Sort.unsorted()), entityGraph).setMaxResults(2).getSingleResult());
+        } catch (NoResultException e) {
+            return Optional.empty();
+        }
+    }
+
+    @Override
+    public List<T> findAll(final Specification<T> spec, final String entityGraph) {
+        return withEntityGraph(getQuery(spec, Sort.unsorted()), entityGraph).getResultList();
+    }
+
+    @Override
+    public Page<T> findAll(final Specification<T> spec, final String entityGraph, final Pageable pageable) {
+        final TypedQuery<T> query = withEntityGraph(getQuery(spec, pageable), entityGraph);
+        return pageable.isUnpaged() ? new PageImpl<>(query.getResultList())
+                : readPage(query, getDomainClass(), pageable, spec);
+    }
+
+    @Override
+    public List<T> findAll(final Specification<T> spec, final String entityGraph, final Sort sort) {
+        return withEntityGraph(getQuery(spec, sort), entityGraph).getResultList();
+    }
+
     @NonNull
     @Override
     public Slice<T> findAllWithoutCount(@Nullable final AccessController.Operation operation, @Nullable Specification<T> spec,
@@ -111,6 +148,11 @@ public String toString() {
         return getClass().getSimpleName() + '<' + getDomainClass().getSimpleName() + '>';
     }
 
+    private TypedQuery<T> withEntityGraph(final TypedQuery<T> query, final String entityGraph) {
+        final EntityGraph<?> graph = ObjectUtils.isEmpty(entityGraph) ? null : entityManager.createEntityGraph(entityGraph);
+        return graph == null ? query : query.setHint("javax.persistence.loadgraph", graph);
+    }
+
     private <S extends T> Page<S> readPageWithoutCount(final TypedQuery<S> query, final Pageable pageable) {
         query.setFirstResult((int) pageable.getOffset());
         query.setMaxResults(pageable.getPageSize());

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/repository/JpaSpecificationEntityGraphExecutor.java

@@ -0,0 +1,81 @@
+/**
+ * Copyright (c) 2024 Contributors to the Eclipse Foundation
+ *
+ * This program and the accompanying materials are made
+ * available under the terms of the Eclipse Public License 2.0
+ * which is available at https://www.eclipse.org/legal/epl-2.0/
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ */
+package org.eclipse.hawkbit.repository.jpa.repository;
+
+import java.util.List;
+import java.util.Optional;
+import java.util.function.Function;
+
+import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.criteria.CriteriaQuery;
+import jakarta.persistence.criteria.Root;
+
+import org.eclipse.hawkbit.repository.model.BaseEntity;
+import org.springframework.data.domain.Page;
+import org.springframework.data.domain.Pageable;
+import org.springframework.data.domain.Slice;
+import org.springframework.data.domain.Sort;
+import org.springframework.data.jpa.domain.Specification;
+import org.springframework.data.repository.query.FluentQuery;
+import org.springframework.lang.Nullable;
+
+/**
+ * Repository interface that offers JpaSpecificationExecutor#findOne/All methods with entity graph loading
+ *
+ * @param <T> entity type
+ */
+public interface JpaSpecificationEntityGraphExecutor<T> {
+
+    /**
+     * Returns a single entity matching the given {@link Specification} with the entity graph hint or {@link Optional#empty()} if none found.
+     * @see org.springframework.data.jpa.repository.JpaSpecificationExecutor#findOne(Specification)
+     *
+     * @param spec must not be {@literal null}.
+     * @param entityGraph the entity graph hint to use
+     * @return never {@literal null}.
+     * @throws org.springframework.dao.IncorrectResultSizeDataAccessException if more than one entity found.
+     */
+    Optional<T> findOne(Specification<T> spec, String entityGraph);
+
+    /**
+     * Returns all entities matching the given {@link Specification} with the entity graph hint.
+     * <p>
+     * @see org.springframework.data.jpa.repository.JpaSpecificationExecutor#findAll(Specification)
+     *
+     * @param spec can be {@literal null}.
+     * @param entityGraph the entity graph hint to use
+     * @return never {@literal null}.
+     */
+    List<T> findAll(@Nullable Specification<T> spec, String entityGraph);
+
+    /**
+     * Returns a {@link Page} of entities matching the given {@link Specification} with the entity graph hint.
+     * <p>
+     * @see org.springframework.data.jpa.repository.JpaSpecificationExecutor#findAll(Specification, Pageable)
+     *
+     * @param spec can be {@literal null}.
+     * @param entityGraph the entity graph hint to use
+     * @param pageable must not be {@literal null}.
+     * @return never {@literal null}.
+     */
+    Page<T> findAll(@Nullable Specification<T> spec, String entityGraph, Pageable pageable);
+
+    /**
+     * Returns all entities matching the given {@link Specification} and {@link Sort} with the entity graph hint.
+     * <p>
+     * @see org.springframework.data.jpa.repository.JpaSpecificationExecutor#findAll(Specification, Sort)
+     *
+     * @param spec can be {@literal null}.
+     * @param entityGraph the entity graph hint to use
+     * @param sort must not be {@literal null}.
+     * @return never {@literal null}.
+     */
+    List<T> findAll(@Nullable Specification<T> spec, String entityGraph, Sort sort);
+}

hawkbit-repository/hawkbit-repository-jpa/src/main/java/org/eclipse/hawkbit/repository/jpa/specifications/ActionSpecifications.java

@@ -78,16 +78,17 @@ public static Specification<JpaAction> byTargetIdsAndActiveAndStatusAndDSNotRequ
 
     /**
      * Returns active actions by target controller that has null or non-null depending on <code>isNull</code> value.
-     * Fetches action's distribution set.
      *
      * @param controllerId controller id
+     * @param isNull if <code>true</code> return with <code>null</code> weight, otherwise with non-<code>null</code>
      * @return the matching action s.
      */
-    public static Specification<JpaAction> byTargetControllerIdAndActive(final String controllerId) {
+    public static Specification<JpaAction> byTargetControllerIdAndActiveAndWeightIsNull(final String controllerId, final boolean isNull) {
         return (root, query, cb) ->
                 cb.and(
-                    cb.equal(root.get(JpaAction_.target).get(JpaTarget_.controllerId), controllerId),
-                    cb.equal(root.get(JpaAction_.active), true));
+                        cb.equal(root.get(JpaAction_.target).get(JpaTarget_.controllerId), controllerId),
+                        cb.equal(root.get(JpaAction_.active), true),
+                        isNull ? cb.isNull(root.get(JpaAction_.weight)) : cb.isNotNull(root.get(JpaAction_.weight)));
     }
 
     public static Specification<JpaAction> byDistributionSetId(final Long distributionSetId) {