Fix non found exception content of DS management (un)assignTag (#1895)

Signed-off-by: Marinov Avgustin <Avgustin.Marinov@bosch.com>

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

@@ -14,9 +14,11 @@ import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.BiFunction;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -73,6 +75,7 @@ import org.eclipse.hawkbit.repository.model.DistributionSetType;
 import org.eclipse.hawkbit.repository.model.MetaData;
 import org.eclipse.hawkbit.repository.model.SoftwareModule;
 import org.eclipse.hawkbit.repository.model.Statistic;
+import org.eclipse.hawkbit.repository.model.Target;
 import org.eclipse.hawkbit.repository.rsql.VirtualPropertyReplacer;
 import org.eclipse.hawkbit.utils.TenantConfigHelper;
 import org.springframework.dao.ConcurrencyFailureException;
@@ -190,45 +193,49 @@ public class JpaDistributionSetManagement implements DistributionSetManagement {
     @Retryable(include = {
             ConcurrencyFailureException.class }, maxAttempts = Constants.TX_RT_MAX, backoff = @Backoff(delay = Constants.TX_RT_DELAY))
     public List<DistributionSet> assignTag(final Collection<Long> ids, final long dsTagId) {
-        return updateTags(
+        return updateTag(ids, dsTagId, (tag, distributionSet) -> {
-                ids,
+            if (distributionSet.getTags().contains(tag)) {
-                () -> distributionSetTagManagement.get(dsTagId).orElseThrow(() -> new EntityNotFoundException(DistributionSetTag.class, dsTagId)),
+                return distributionSet;
-                (allDs, distributionSetTag) -> {
+            } else {
-                    allDs.forEach(ds -> ds.addTag(distributionSetTag));
+                distributionSet.addTag(tag);
-                    return Collections.unmodifiableList(distributionSetRepository.saveAll(allDs));
+                return distributionSetRepository.save(distributionSet);
-                });
+            }
+        });
     }
-
     @Override
     @Transactional
     @Retryable(include = {
             ConcurrencyFailureException.class }, maxAttempts = Constants.TX_RT_MAX, backoff = @Backoff(delay = Constants.TX_RT_DELAY))
     public List<DistributionSet> unassignTag(final Collection<Long> ids, final long dsTagId) {
-        return updateTags(
+        return updateTag(ids, dsTagId, (tag, distributionSet) -> {
-                ids,
+                    if (distributionSet.getTags().contains(tag)) {
-                () -> distributionSetTagManagement.get(dsTagId).orElseThrow(() -> new EntityNotFoundException(DistributionSetTag.class, dsTagId)),
+                        distributionSet.removeTag(tag);
-                (allDs, distributionSetTag) -> {
+                        return distributionSetRepository.save(distributionSet);
-                    allDs.forEach(ds -> ds.removeTag(distributionSetTag));
+                    } else {
-                    return Collections.unmodifiableList(distributionSetRepository.saveAll(allDs));
+                        return distributionSet;
+                    }
                 });
     }
-
+    private List<DistributionSet> updateTag(
-    private <T> T updateTags(
+            final Collection<Long> dsIds, final long dsTagId,
-            final Collection<Long> dsIds, final Supplier<DistributionSetTag> tagSupplier,
+            final BiFunction<DistributionSetTag, JpaDistributionSet, DistributionSet> updater) {
-            final BiFunction<List<JpaDistributionSet>, DistributionSetTag, T> updater) {
+        final DistributionSetTag tag = distributionSetTagManagement.get(dsTagId)
+                .orElseThrow(() -> new EntityNotFoundException(DistributionSetTag.class, dsTagId));
         final List<JpaDistributionSet> allDs =  dsIds.size() == 1 ?
-                distributionSetRepository.findById(dsIds.iterator().next()).map(List::of).orElseGet(Collections::emptyList) :
+                distributionSetRepository.findById(dsIds.iterator().next())
+                        .map(List::of)
+                        .orElseGet(Collections::emptyList) :
                 distributionSetRepository.findAll(DistributionSetSpecification.byIds(dsIds));
         if (allDs.size() < dsIds.size()) {
-            throw new EntityNotFoundException(DistributionSet.class, dsIds, allDs.stream().map(DistributionSet::getId).toList());
+            throw new EntityNotFoundException(DistributionSet.class, notFound(dsIds, allDs));
         }
 
-        final DistributionSetTag distributionSetTag = tagSupplier.get();
         try {
-            return updater.apply(allDs, distributionSetTag);
+            // apply update and collect modified targets
+            return allDs.stream().map(distributionSet -> updater.apply(tag, distributionSet)).toList();
         } finally {
             // No reason to save the tag
-            entityManager.detach(distributionSetTag);
+            entityManager.detach(tag);
         }
     }
 
@@ -791,6 +798,13 @@ public class JpaDistributionSetManagement implements DistributionSetManagement {
         return true;
     }
 
+    private static Collection<Long> notFound(final Collection<Long> distributionSetIds, final List<JpaDistributionSet> foundDistributionSets) {
+        final Map<Long, JpaDistributionSet> foundDistributionSetMap = foundDistributionSets.stream()
+                .collect(Collectors.toMap(JpaDistributionSet::getId, Function.identity()));
+        return distributionSetIds.stream().filter(id -> !foundDistributionSetMap.containsKey(id)).toList();
+    }
+
+
     private void lockSoftwareModules(final JpaDistributionSet distributionSet) {
         distributionSet.getModules().forEach(module -> {
             if (!module.isLocked()) {
@@ -824,7 +838,7 @@ public class JpaDistributionSetManagement implements DistributionSetManagement {
     @Retryable(include = {
             ConcurrencyFailureException.class }, maxAttempts = Constants.TX_RT_MAX, backoff = @Backoff(delay = Constants.TX_RT_DELAY))
     public DistributionSetTagAssignmentResult toggleTagAssignment(final Collection<Long> ids, final String tagName) {
-        return updateTags(
+        return updateTag(
                 ids,
                 () -> distributionSetTagManagement
                         .getByName(tagName)
@@ -855,6 +869,24 @@ public class JpaDistributionSetManagement implements DistributionSetManagement {
                     return result;
                 });
     }
+    private <T> T updateTag(
+            final Collection<Long> dsIds, final Supplier<DistributionSetTag> tagSupplier,
+            final BiFunction<List<JpaDistributionSet>, DistributionSetTag, T> updater) {
+        final List<JpaDistributionSet> allDs =  dsIds.size() == 1 ?
+                distributionSetRepository.findById(dsIds.iterator().next()).map(List::of).orElseGet(Collections::emptyList) :
+                distributionSetRepository.findAll(DistributionSetSpecification.byIds(dsIds));
+        if (allDs.size() < dsIds.size()) {
+            throw new EntityNotFoundException(DistributionSet.class, dsIds, allDs.stream().map(DistributionSet::getId).toList());
+        }
+
+        final DistributionSetTag distributionSetTag = tagSupplier.get();
+        try {
+            return updater.apply(allDs, distributionSetTag);
+        } finally {
+            // No reason to save the tag
+            entityManager.detach(distributionSetTag);
+        }
+    }
 
     @Override
     @Transactional

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

@@ -559,22 +559,29 @@ public class JpaTargetManagement implements TargetManagement {
             }
         });
     }
-    private List<Target> updateTag(final Collection<String> controllerIds, final long targetTagId, final BiFunction<JpaTargetTag, JpaTarget, Target> updater) {
+    private List<Target> updateTag(
+            final Collection<String> controllerIds, final long targetTagId,
+            final BiFunction<JpaTargetTag, JpaTarget, Target> updater) {
         final JpaTargetTag tag = targetTagRepository.findById(targetTagId)
                 .orElseThrow(() -> new EntityNotFoundException(TargetTag.class, targetTagId));
-        final List<JpaTarget> targets = targetRepository
+        final List<JpaTarget> targets = controllerIds.size() == 1 ?
-                .findAll(TargetSpecifications.byControllerIdWithTagsInJoin(controllerIds));
+                targetRepository.findOne(TargetSpecifications.hasControllerId(controllerIds.iterator().next()))
+                        .map(List::of)
+                        .orElseGet(Collections::emptyList) :
+                targetRepository
+                        .findAll(TargetSpecifications.byControllerIdWithTagsInJoin(controllerIds));
         if (targets.size() < controllerIds.size()) {
             throw new EntityNotFoundException(Target.class, notFound(controllerIds, targets));
         }
         targetRepository.getAccessController()
                 .ifPresent(acm -> acm.assertOperationAllowed(AccessController.Operation.UPDATE, targets));
 
-        // apply update and collect modified targets
+        try {
-        final List<Target> result = targets.stream().map(target -> updater.apply(tag, target)).toList();
+            return targets.stream().map(target -> updater.apply(tag, target)).toList();
-        // No reason to save the tag
+        } finally {
-        entityManager.detach(tag);
+            // No reason to save the tag
-        return result;
+            entityManager.detach(tag);
+        }
     }
 
     @Override