Return immutable result from TaskStore::query

Author: FlorentinD

core/src/main/java/org/neo4j/gds/core/utils/progress/EmptyTaskStore.java

@@ -46,8 +46,8 @@ public Stream<UserTask> query(JobId jobId) {
     }
 
     @Override
-    public @NotNull Map<JobId, Task> query(String username) {
-        return Map.of();
+    public @NotNull Stream<UserTask> query(String username) {
+        return Stream.empty();
     }
 
     @Override

core/src/main/java/org/neo4j/gds/core/utils/progress/GlobalTaskStore.java

@@ -72,8 +72,12 @@ public Stream<UserTask> query(JobId jobId) {
     }
 
     @Override
-    public @NotNull Map<JobId, Task> query(String username) {
-        return registeredTasks.getOrDefault(username, Map.of());
+    public Stream<UserTask> query(String username) {
+        return registeredTasks
+            .getOrDefault(username, Map.of())
+            .entrySet()
+            .stream()
+            .map(jobIdTask -> ImmutableUserTask.of(username, jobIdTask.getKey(), jobIdTask.getValue()));
     }
 
     @Override

core/src/main/java/org/neo4j/gds/core/utils/progress/TaskRegistry.java

@@ -55,7 +55,7 @@ public void unregisterTask() {
 
     public boolean containsTask(Task task) {
         return taskStore.query(username, jobId)
-            .map(t -> t == task)
+            .map(t -> t.task() == task)
             .orElse(false);
     }
 }

core/src/main/java/org/neo4j/gds/core/utils/progress/TaskStore.java

@@ -37,8 +37,7 @@ public interface TaskStore {
 
     Stream<UserTask> query(JobId jobId);
 
-    @NotNull
-    Map<JobId, Task> query(String username);
+    Stream<UserTask> query(String username);
 
     Optional<UserTask> query(String username, JobId jobId);
 

core/src/test/java/org/neo4j/gds/core/utils/progress/BaseTaskRegistryExtensionTest.java

@@ -33,6 +33,7 @@
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toList;
+import static org.neo4j.gds.core.utils.progress.TaskStore.UserTask;
 
 abstract class BaseTaskRegistryExtensionTest extends BaseTest {
     abstract boolean featureEnabled();
@@ -79,7 +80,7 @@ public static class ProgressProc {
 
         @Procedure("gds.test.log")
         public Stream<Bar> foo() {
-            return taskStore.query("").values().stream().map(Task::description).map(Bar::new);
+            return taskStore.query("").map(UserTask::task).map(Task::description).map(Bar::new);
         }
     }
 

core/src/test/java/org/neo4j/gds/core/utils/progress/LocalTaskRegistryFactoryTest.java

@@ -53,7 +53,7 @@ void shouldPutAndRemoveDistinctTasks() {
 
         taskRegistry1.unregisterTask();
 
-        assertThat(taskStore.query("")).containsValue(task2).doesNotContainValue(task1);
+        assertThat(taskStore.query("").map(TaskStore.UserTask::task)).contains(task2).doesNotContain(task1);
     }
 
     @Test

core/src/test/java/org/neo4j/gds/core/utils/progress/TaskRegistryTest.java

@@ -20,6 +20,7 @@
 package org.neo4j.gds.core.utils.progress;
 
 import org.junit.jupiter.api.Test;
+import org.neo4j.gds.core.utils.progress.TaskStore.UserTask;
 import org.neo4j.gds.core.utils.progress.tasks.Tasks;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -36,14 +37,14 @@ void shouldStoreIncomingTasks() {
         var task1 = Tasks.leaf("task1");
         taskRegistry1.registerTask(task1);
 
-        assertThat(globalTaskStore.query("")).containsValue(task1);
+        assertThat(globalTaskStore.query("").map(UserTask::task)).contains(task1);
         assertThat(globalTaskStore.isEmpty()).isFalse();
 
         var taskRegistry2 = new TaskRegistry("", globalTaskStore);
         var task2 = Tasks.leaf("task2");
         taskRegistry2.registerTask(task2);
 
-        assertThat(globalTaskStore.query("")).containsValues(task1, task2);
+        assertThat(globalTaskStore.query("").map(UserTask::task)).contains(task1, task2);
         assertThat(globalTaskStore.isEmpty()).isFalse();
     }
 
@@ -57,7 +58,7 @@ void shouldRemoveStoredTasks() {
 
         assertThat(globalTaskStore.isEmpty()).isFalse();
 
-        var jobId = globalTaskStore.query("").keySet().iterator().next();
+        var jobId = globalTaskStore.query("").map(UserTask::jobId).iterator().next();
         globalTaskStore.remove("", jobId);
 
         assertThat(globalTaskStore.isEmpty()).isTrue();

core/src/test/java/org/neo4j/gds/core/utils/progress/tasks/TaskProgressTrackerTest.java

@@ -26,6 +26,7 @@
 import org.neo4j.gds.core.utils.progress.EmptyTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.GlobalTaskStore;
 import org.neo4j.gds.core.utils.progress.TaskRegistry;
+import org.neo4j.gds.core.utils.progress.TaskStore;
 import org.neo4j.logging.Log;
 
 import java.util.List;
@@ -205,7 +206,7 @@ void shouldRegisterBaseTaskOnBaseTaskStart() {
 
         progressTracker.beginSubTask();
 
-        assertThat(taskStore.query("")).containsValue(task);
+        assertThat(taskStore.query("").map(TaskStore.UserTask::task)).contains(task);
     }
 
     @Test

proc/misc/src/main/java/org/neo4j/gds/ListProgressProc.java

@@ -64,11 +64,7 @@ private Stream<ProgressResult> jobsSummaryView() {
         if (isGdsAdmin()) {
             return taskStore.query().map(ProgressResult::fromTaskStoreEntry);
         } else {
-            return taskStore
-                .query(username())
-                .entrySet()
-                .stream()
-                .map(entry -> ProgressResult.fromTaskStoreEntry(username(), entry));
+            return taskStore.query(username()).map(ProgressResult::fromTaskStoreEntry);
         }
     }