From 136b9c3e6eada15ed5d8c6e26e02208231a74b25 Mon Sep 17 00:00:00 2001
From: Francis Li <mail@francisli.com>
Date: Thu, 23 Jan 2025 18:24:25 -0800
Subject: [PATCH] Update assignTasks implementation

---
 src/backend/routers/iep.test.ts               | 21 ++++---
 src/backend/routers/iep.ts                    | 62 +++++++++++++------
 .../benchmarks/BenchmarkAssignmentModal.tsx   | 32 +++++-----
 3 files changed, 71 insertions(+), 44 deletions(-)

diff --git a/src/backend/routers/iep.test.ts b/src/backend/routers/iep.test.ts
index d27e66ad..f0e7dd53 100644
--- a/src/backend/routers/iep.test.ts
+++ b/src/backend/routers/iep.test.ts
@@ -165,7 +165,7 @@ test("addTask - no duplicate benchmark_id + assigned_id combo", async (t) => {
 });
 
 test("assignTaskToParas - no duplicate benchmark_id + para_id combo", async (t) => {
-  const { trpc, seed } = await getTestServer(t, {
+  const { trpc, db, seed } = await getTestServer(t, {
     authenticateAs: UserType.CaseManager,
   });
 
@@ -212,17 +212,18 @@ test("assignTaskToParas - no duplicate benchmark_id + para_id combo", async (t)
     para_ids: [para_1.user_id],
   });
 
-  const error = await t.throwsAsync(async () => {
-    await trpc.iep.assignTaskToParas.mutate({
-      benchmark_id: benchmark1Id,
-      para_ids: [para_1.user_id, para_2.user_id],
-    });
+  await trpc.iep.assignTaskToParas.mutate({
+    benchmark_id: benchmark1Id,
+    para_ids: [para_1.user_id, para_2.user_id],
   });
 
-  t.is(
-    error?.message,
-    "Task already exists: This benchmark has already been assigned to one or more of these paras"
-  );
+  const result = await db
+    .selectFrom("task")
+    .where("benchmark_id", "=", benchmark1Id)
+    .selectAll()
+    .execute();
+
+  t.is(result.length, 2);
 });
 
 test("add benchmark - check full schema", async (t) => {
diff --git a/src/backend/routers/iep.ts b/src/backend/routers/iep.ts
index 2d82bc8f..01aecee5 100644
--- a/src/backend/routers/iep.ts
+++ b/src/backend/routers/iep.ts
@@ -237,32 +237,56 @@ export const iep = router({
     .mutation(async (req) => {
       const { benchmark_id, para_ids } = req.input;
 
-      // TODO: instead of throwing an error if a task already exists,
-      // just add/remove existing tasks as needed.
-
+      // fetch all existing task records for this benchmark
       const existingTasks = await req.ctx.db
         .selectFrom("task")
         .where("benchmark_id", "=", benchmark_id)
-        .where("assignee_id", "in", para_ids)
         .selectAll()
         .execute();
 
-      if (existingTasks.length > 0) {
-        throw new Error(
-          "Task already exists: This benchmark has already been assigned to one or more of these paras"
-        );
-      }
-
-      const result = await req.ctx.db
-        .insertInto("task")
-        .values(
-          para_ids.map((para_id) => ({
-            benchmark_id,
-            assignee_id: para_id,
-          }))
+      // collect list of task records to delete (if assignee_id not in para_ids)
+      const deleteTaskIds = existingTasks
+        .filter(
+          (task) => task.assignee_id && !para_ids.includes(task.assignee_id)
         )
-        .returningAll()
-        .executeTakeFirst();
+        .map((task) => task.task_id);
+
+      // collect a list of new assignee_ids to insert
+      const newAssigneeIds = para_ids.filter(
+        (para_id) => !existingTasks.find((task) => task.assignee_id === para_id)
+      );
+
+      let result:
+        | {
+            benchmark_id: string | null;
+            assignee_id: string | null;
+            created_at: Date;
+            seen: boolean;
+            task_id: string;
+          }
+        | undefined;
+      await req.ctx.db.transaction().execute(async (trx) => {
+        // delete tasks whose assignees that are not in para_ids
+        if (deleteTaskIds.length > 0) {
+          await trx
+            .deleteFrom("task")
+            .where("task_id", "in", deleteTaskIds)
+            .execute();
+        }
+        // insert new tasks for new assignees
+        if (newAssigneeIds.length > 0) {
+          result = await req.ctx.db
+            .insertInto("task")
+            .values(
+              newAssigneeIds.map((para_id) => ({
+                benchmark_id,
+                assignee_id: para_id,
+              }))
+            )
+            .returningAll()
+            .executeTakeFirst();
+        }
+      });
       return result;
     }),
 
diff --git a/src/components/benchmarks/BenchmarkAssignmentModal.tsx b/src/components/benchmarks/BenchmarkAssignmentModal.tsx
index ddfd3a83..bde38865 100644
--- a/src/components/benchmarks/BenchmarkAssignmentModal.tsx
+++ b/src/components/benchmarks/BenchmarkAssignmentModal.tsx
@@ -97,21 +97,23 @@ export const BenchmarkAssignmentModal = (
       // Reached end, save
       try {
         // maybe invoke these two in parallel?
-        await assignTaskToPara.mutateAsync({
-          benchmark_id: props.benchmark_id,
-          para_ids: selectedParaIds,
-        });
-        await updateBenchmark.mutateAsync({
-          benchmark_id: props.benchmark_id,
-          due_date:
-            assignmentDuration.type === "until_date"
-              ? assignmentDuration.date
-              : undefined,
-          trial_count:
-            assignmentDuration.type === "minimum_number_of_collections"
-              ? assignmentDuration.minimumNumberOfCollections
-              : undefined,
-        });
+        await Promise.all([
+          assignTaskToPara.mutateAsync({
+            benchmark_id: props.benchmark_id,
+            para_ids: selectedParaIds,
+          }),
+          updateBenchmark.mutateAsync({
+            benchmark_id: props.benchmark_id,
+            due_date:
+              assignmentDuration.type === "until_date"
+                ? assignmentDuration.date
+                : undefined,
+            trial_count:
+              assignmentDuration.type === "minimum_number_of_collections"
+                ? assignmentDuration.minimumNumberOfCollections
+                : undefined,
+          }),
+        ]);
         handleClose();
       } catch (err) {
         // TODO: issue #450