From 0a01a210904ae8d6649162a58718d7b24bbacc37 Mon Sep 17 00:00:00 2001
From: eagle <aparokshambhargav@gmail.com>
Date: Fri, 23 Aug 2024 15:08:26 +0530
Subject: [PATCH] fix: ballot issue

---
 src/features/ballot/types/index.ts            |  5 +++++
 src/features/voters/hooks/useApproveVoters.ts | 14 +++++++++++++-
 src/server/api/routers/ballot.ts              | 19 +++++++++++++++++++
 src/server/api/trpc.ts                        | 16 ++++++++--------
 4 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/src/features/ballot/types/index.ts b/src/features/ballot/types/index.ts
index 4ad5cef0..8dbab02e 100644
--- a/src/features/ballot/types/index.ts
+++ b/src/features/ballot/types/index.ts
@@ -28,6 +28,11 @@ export const AllocationSchema = z.object({
   amount: z.number().min(0),
   locked: z.boolean().default(true),
 });
+export const BallotV2Schema1 = z.object({
+  voterId: z.string(),
+  roundId: z.string(),
+  type: z.string(),
+});
 export const BallotV2Schema = z.object({
   allocations: z.array(AllocationSchema),
 });
diff --git a/src/features/voters/hooks/useApproveVoters.ts b/src/features/voters/hooks/useApproveVoters.ts
index f55e253a..c607fdb7 100644
--- a/src/features/voters/hooks/useApproveVoters.ts
+++ b/src/features/voters/hooks/useApproveVoters.ts
@@ -25,6 +25,8 @@ export function useApproveVoters({
   const signer = useEthersSigner();
   const { data: round } = useCurrentRound();
 
+  const saveBallot = api.ballot.saveBallot.useMutation();
+
   return useMutation({
     mutationFn: async (voters: string[]) => {
       if (!signer) throw new Error("Connect wallet first");
@@ -45,9 +47,19 @@ export function useApproveVoters({
           ),
         ),
       );
-      return attest.mutateAsync(
+      await attest.mutateAsync(
         attestations.map((att) => ({ ...att, data: [att.data] })),
       );
+
+      await Promise.all(
+        voters.map((voter) =>
+          saveBallot.mutateAsync({
+            voterId: voter,
+            roundId: round.id,
+            type: round.type,
+          }),
+        ),
+      );
     },
     onSuccess,
     onError,
diff --git a/src/server/api/routers/ballot.ts b/src/server/api/routers/ballot.ts
index dc6fe302..c98fbbab 100644
--- a/src/server/api/routers/ballot.ts
+++ b/src/server/api/routers/ballot.ts
@@ -10,6 +10,7 @@ import {
   AllocationSchema,
   Allocation,
   BallotPublishSchema,
+  BallotV2Schema1,
 } from "~/features/ballot/types";
 import { calculateBalancedAllocations } from "~/features/ballot/hooks/useBallotEditor";
 import { TRPCError } from "@trpc/server";
@@ -43,6 +44,24 @@ export const ballotRouter = createTRPCRouter({
       select: defaultBallotSelect,
     });
   }),
+  saveBallot: ballotProcedure
+    .input(BallotV2Schema1)
+    .mutation(async ({ input, ctx }) => {
+      const voterId = input.voterId;
+      const roundId = input.roundId;
+      const type = input.type as RoundTypes;
+
+      return ctx.db.ballotV2.upsert({
+        where: { voterId_roundId_type: { voterId, roundId, type } },
+        update: {},
+        create: {
+          voterId,
+          roundId,
+          type,
+        },
+        select: defaultBallotSelect,
+      });
+    }),
   save: ballotProcedure
     .input(AllocationSchema)
     .mutation(async ({ input, ctx }) => {
diff --git a/src/server/api/trpc.ts b/src/server/api/trpc.ts
index 724a59af..106e0512 100644
--- a/src/server/api/trpc.ts
+++ b/src/server/api/trpc.ts
@@ -251,15 +251,15 @@ export const ballotProcedure = protectedProcedure.use(roundMiddleware).use(
   t.middleware(async ({ ctx, next }) => {
     const voterId = ctx.session?.user.name!;
     const roundId = ctx.round?.id!;
-    const type = ctx.round?.type as RoundTypes;
-
-    // Find or create ballot
-    const ballot = await ctx.db.ballotV2.upsert({
-      where: { voterId_roundId_type: { voterId, roundId, type } },
-      update: {},
-      create: { voterId, roundId, type },
-      include: { allocations: true },
+
+    // Find ballot
+    const ballot = await ctx.db.ballotV2.findFirst({
+      where: { voterId, roundId },
+      select: { id: true },
     });
+    if (!ballot) {
+      throw new TRPCError({ code: "NOT_FOUND", message: "Ballot not found" });
+    }
     return next({
       ctx: { ...ctx, voterId, roundId, ballotId: ballot.id, ballot },
     });