Allow querying TaskStore over all users

FlorentinD · FlorentinD · commit 7d271208465e · 2022-11-28T11:45:55.000+01:00
This is intended for admins.
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/EmptyTaskStore.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/EmptyTaskStore.java
@@ -35,13 +35,23 @@ public void store(String username, JobId jobId, Task task) {}
     @Override
     public void remove(String username, JobId jobId) {}
 
+    @Override
+    public Stream<UserTask> query() {
+        return Stream.empty();
+    }
+
+    @Override
+    public Stream<UserTask> query(JobId jobId) {
+        return Stream.empty();
+    }
+
     @Override
     public @NotNull Map<JobId, Task> query(String username) {
         return Map.of();
     }
 
     @Override
-    public Optional<Task> query(String username, JobId jobId) {
+    public Optional<UserTask> query(String username, JobId jobId) {
         return Optional.empty();
     }
 
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/GlobalTaskStore.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/GlobalTaskStore.java
@@ -54,15 +54,34 @@ public void remove(String username, JobId jobId) {
         }
     }
 
+    @Override
+    public Stream<UserTask> query() {
+        return registeredTasks
+            .entrySet()
+            .stream()
+            .flatMap(tasksPerUsers -> tasksPerUsers
+                .getValue()
+                .entrySet()
+                .stream()
+                .map(jobTask -> ImmutableUserTask.of(tasksPerUsers.getKey(), jobTask.getKey(), jobTask.getValue())));
+    }
+
+    @Override
+    public Stream<UserTask> query(JobId jobId) {
+        return query().filter(userTask -> userTask.jobId().equals(jobId));
+    }
+
     @Override
     public @NotNull Map<JobId, Task> query(String username) {
         return registeredTasks.getOrDefault(username, Map.of());
     }
 
     @Override
-    public Optional<Task> query(String username, JobId jobId) {
+    public Optional<UserTask> query(String username, JobId jobId) {
         if (registeredTasks.containsKey(username)) {
-            return Optional.ofNullable(registeredTasks.get(username).get(jobId));
+            return Optional
+                .ofNullable(registeredTasks.get(username).get(jobId))
+                .map(task -> ImmutableUserTask.of(username, jobId, task));
         }
         return Optional.empty();
     }
diff --git a/core/src/main/java/org/neo4j/gds/core/utils/progress/TaskStore.java b/core/src/main/java/org/neo4j/gds/core/utils/progress/TaskStore.java
@@ -20,6 +20,7 @@
 package org.neo4j.gds.core.utils.progress;
 
 import org.jetbrains.annotations.NotNull;
+import org.neo4j.gds.annotation.ValueClass;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 
 import java.util.Map;
@@ -32,14 +33,27 @@ public interface TaskStore {
 
     void remove(String username, JobId jobId);
 
+    Stream<UserTask> query();
+
+    Stream<UserTask> query(JobId jobId);
+
     @NotNull
     Map<JobId, Task> query(String username);
 
-    Optional<Task> query(String username, JobId jobId);
+    Optional<UserTask> query(String username, JobId jobId);
 
     Stream<Task> taskStream();
 
     boolean isEmpty();
 
     long taskCount();
+
+    @ValueClass
+    interface UserTask {
+        String username();
+
+        JobId jobId();
+
+        Task task();
+    }
 }
diff --git a/core/src/test/java/org/neo4j/gds/core/utils/progress/GlobalTaskStoreTest.java b/core/src/test/java/org/neo4j/gds/core/utils/progress/GlobalTaskStoreTest.java
@@ -58,4 +58,15 @@ void shouldCountAcrossUsers() {
 
         assertThat(taskStore.taskCount()).isEqualTo(3);
     }
+
+    @Test
+    void shouldQueryMultipleUsers() {
+        var taskStore = new GlobalTaskStore();
+        taskStore.store("alice", new JobId("42"), Tasks.leaf("leaf"));
+        taskStore.store("bob", new JobId("1337"), Tasks.leaf("other"));
+
+        assertThat(taskStore.query()).hasSize(2);
+        assertThat(taskStore.query(new JobId("42"))).hasSize(1);
+        assertThat(taskStore.query(new JobId(""))).hasSize(0);
+    }
 }
diff --git a/proc/misc/src/main/java/org/neo4j/gds/ListProgressProc.java b/proc/misc/src/main/java/org/neo4j/gds/ListProgressProc.java
@@ -38,6 +38,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import static org.neo4j.gds.utils.StringFormatting.formatWithLocale;
@@ -60,21 +61,53 @@ public Stream<ProgressResult> listProgress(
     }
 
     private Stream<ProgressResult> jobsSummaryView() {
-        return taskStore.query(username()).entrySet().stream().map(ProgressResult::fromTaskStoreEntry);
+        if (isGdsAdmin()) {
+            return taskStore.query().map(ProgressResult::fromTaskStoreEntry);
+        } else {
+            return taskStore
+                .query(username())
+                .entrySet()
+                .stream()
+                .map(entry -> ProgressResult.fromTaskStoreEntry(username(), entry));
+        }
     }
 
     private Stream<ProgressResult> jobDetailView(String jobIdAsString) {
         var jobId = new JobId(jobIdAsString);
-        var task = taskStore.query(username(), jobId).orElseThrow(
-            () -> new IllegalArgumentException(formatWithLocale("No task with job id `%s` was found.", jobIdAsString))
-        );
-        var jobProgressVisitor = new JobProgressVisitor(jobId);
-        TaskTraversal.visitPreOrderWithDepth(task, jobProgressVisitor);
+
+        if (isGdsAdmin()) {
+            var progressResults = taskStore
+                .query(jobId)
+                .flatMap(ListProgressProc::jobProgress)
+                .collect(Collectors.toList());
+
+            if (progressResults.isEmpty()) {
+                throw new IllegalArgumentException(formatWithLocale(
+                    "No task with job id `%s` was found.",
+                    jobIdAsString
+                ));
+            }
+
+            return progressResults.stream();
+        } else {
+            return taskStore.query(username(), jobId).map(ListProgressProc::jobProgress).orElseThrow(
+                () -> new IllegalArgumentException(formatWithLocale(
+                    "No task with job id `%s` was found.",
+                    jobIdAsString
+                ))
+            );
+        }
+    }
+
+    private static Stream<ProgressResult> jobProgress(TaskStore.UserTask userTask) {
+        var jobProgressVisitor = new JobProgressVisitor(userTask.jobId(), userTask.username());
+        TaskTraversal.visitPreOrderWithDepth(userTask.task(), jobProgressVisitor);
         return jobProgressVisitor.progressRowsStream();
     }
 
     @SuppressWarnings("unused")
     public static class ProgressResult {
+        public String username;
         public String jobId;
         public String taskName;
         public String progress;
@@ -83,22 +116,27 @@ public static class ProgressResult {
         public LocalTimeValue timeStarted;
         public String elapsedTime;
 
-        static ProgressResult fromTaskStoreEntry(Map.Entry<JobId, Task> taskStoreEntry) {
+        static ProgressResult fromTaskStoreEntry(String username, Map.Entry<JobId, Task> taskStoreEntry) {
             var jobId = taskStoreEntry.getKey();
             var task = taskStoreEntry.getValue();
-            return new ProgressResult(task, jobId, task.description());
+            return new ProgressResult(username, task, jobId, task.description());
+        }
+
+        static ProgressResult fromTaskStoreEntry(TaskStore.UserTask userTask) {
+            return new ProgressResult(userTask.username(), userTask.task(), userTask.jobId(), userTask.task().description());
         }
 
-        static ProgressResult fromTaskWithDepth(Task task, JobId jobId, int depth) {
+        static ProgressResult fromTaskWithDepth(String username, Task task, JobId jobId, int depth) {
             var treeViewTaskName = StructuredOutputHelper.treeViewDescription(task.description(), depth);
-            return new ProgressResult(task, jobId, treeViewTaskName);
+            return new ProgressResult(username, task, jobId, treeViewTaskName);
         }
 
-        public ProgressResult(Task task, JobId jobId, String taskName) {
+        public ProgressResult(String username, Task task, JobId jobId, String taskName) {
             var progressContainer = task.getProgress();
 
             this.jobId = jobId.asString();
             this.taskName = taskName;
+            this.username = username;
             this.progress = StructuredOutputHelper.computeProgress(progressContainer);
             this.progressBar = StructuredOutputHelper.progressBar(progressContainer, PROGRESS_BAR_LENGTH);
             this.status = task.status().name();
@@ -132,10 +170,12 @@ private String prettyElapsedTime(Task task) {
     public static class JobProgressVisitor extends DepthAwareTaskVisitor {
 
         private final JobId jobId;
+        private final String username;
         private final List<ProgressResult> progressRows;
 
-        JobProgressVisitor(JobId jobId) {
+        JobProgressVisitor(JobId jobId, String username) {
             this.jobId = jobId;
+            this.username = username;
             this.progressRows = new ArrayList<>();
         }
 
@@ -145,7 +185,7 @@ Stream<ProgressResult> progressRowsStream() {
 
         @Override
         public void visit(Task task) {
-            progressRows.add(ProgressResult.fromTaskWithDepth(task, jobId, depth()));
+            progressRows.add(ProgressResult.fromTaskWithDepth(username, task, jobId, depth()));
         }
     }
 }
diff --git a/proc/misc/src/test/java/org/neo4j/gds/ListProgressProcTest.java b/proc/misc/src/test/java/org/neo4j/gds/ListProgressProcTest.java
@@ -23,6 +23,7 @@
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.beta.generator.GraphGenerateProc;
 import org.neo4j.gds.compat.GraphDatabaseApiProxy;
+import org.neo4j.gds.core.Username;
 import org.neo4j.gds.core.utils.RenamesCurrentThread;
 import org.neo4j.gds.core.utils.progress.JobId;
 import org.neo4j.gds.embeddings.fastrp.FastRP;
@@ -72,11 +73,12 @@ void canListProgressEvent() {
         runQuery("CALL gds.test.pl('foo')");
         assertCypherResult(
             "CALL gds.beta.listProgress() " +
-            "YIELD taskName, progress, progressBar, status, timeStarted, elapsedTime " +
-            "RETURN taskName, progress, progressBar, status, timeStarted, elapsedTime ",
+            "YIELD username, taskName, progress, progressBar, status, timeStarted, elapsedTime " +
+            "RETURN username, taskName, progress, progressBar, status, timeStarted, elapsedTime ",
             List.of(
                 Map.of(
                     "taskName","foo",
+                    "username", Username.EMPTY_USERNAME.username(),
                     "progress", "33.33%",
                     "progressBar", "[###~~~~~~~]",
                     "status", "RUNNING",
diff --git a/test-utils/src/main/java/org/neo4j/gds/InspectableTestProgressTracker.java b/test-utils/src/main/java/org/neo4j/gds/InspectableTestProgressTracker.java
@@ -77,13 +77,13 @@ public void logProgress(long progress) {
     @Override
     public void beginSubTask() {
         super.beginSubTask();
-        progressHistory.add(taskStore.query(userName, jobId).map(Task::getProgress));
+        progressHistory.add(taskStore.query(userName, jobId).map(userTask -> userTask.task().getProgress()));
     }
 
     @Override
     public void endSubTask() {
         super.endSubTask();
-        progressHistory.add(taskStore.query(userName, jobId).map(Task::getProgress));
+        progressHistory.add(taskStore.query(userName, jobId).map(userTask -> userTask.task().getProgress()));
     }
 
     @Override

Original file line number	Diff line number	Diff line change
`@@ -77,13 +77,13 @@ public void logProgress(long progress) {`
`77`	`77`	`@Override`
`78`	`78`	`public void beginSubTask() {`
`79`	`79`	`super.beginSubTask();`
`80`		`- progressHistory.add(taskStore.query(userName, jobId).map(Task::getProgress));`
	`80`	`+ progressHistory.add(taskStore.query(userName, jobId).map(userTask -> userTask.task().getProgress()));`
`81`	`81`	`}`
`82`	`82`
`83`	`83`	`@Override`
`84`	`84`	`public void endSubTask() {`
`85`	`85`	`super.endSubTask();`
`86`		`- progressHistory.add(taskStore.query(userName, jobId).map(Task::getProgress));`
	`86`	`+ progressHistory.add(taskStore.query(userName, jobId).map(userTask -> userTask.task().getProgress()));`
`87`	`87`	`}`
`88`	`88`
`89`	`89`	`@Override`