Add poll database schema and basic implementation

Co-authored-by: CPlusPatch <42910258+CPlusPatch@users.noreply.github.com>

classes/database/note.ts

@@ -28,6 +28,7 @@ import { mergeAndDeduplicate } from "@/lib.ts";
 import { sanitizedHtmlStrip } from "@/sanitization";
 import { contentToHtml, findManyNotes } from "~/classes/functions/status";
 import { config } from "~/config.ts";
+import { Poll } from "./poll.ts";
 import * as VersiaEntities from "~/packages/sdk/entities/index.ts";
 import type { NonTextContentFormatSchema } from "~/packages/sdk/schemas/contentformat.ts";
 import { DeliveryJobType, deliveryQueue } from "../queues/delivery.ts";
@@ -55,6 +56,7 @@ type NoteTypeWithRelations = NoteType & {
     muted: boolean;
     liked: boolean;
     reactions: Omit<typeof Reaction.$type, "note" | "author">[];
+    poll: typeof Poll.$type | null;
 };
 
 export type NoteTypeWithoutRecursiveRelations = Omit<
@@ -691,8 +693,7 @@ export class Note extends BaseInterface<typeof Notes, NoteTypeWithRelations> {
             language: null,
             muted: data.muted,
             pinned: data.pinned,
-            // TODO: Add polls
-            poll: null,
+            poll: data.poll ? data.poll.toApi(userFetching) : null,
             reblog: data.reblog
                 ? await new Note(data.reblog as NoteTypeWithRelations).toApi(
                       userFetching,

classes/database/poll.ts

@@ -0,0 +1,284 @@
+import { db } from "@versia/kit/db";
+import {
+    Notes,
+    PollOptions,
+    Polls,
+    PollVotes,
+    type Users,
+} from "@versia/kit/tables";
+import {
+    and,
+    eq,
+    type InferInsertModel,
+    type InferSelectModel,
+    inArray,
+} from "drizzle-orm";
+import type { z } from "zod";
+import type { Poll as PollSchema } from "@versia/client/schemas";
+import { BaseInterface } from "./base.ts";
+
+/**
+ * Type definition for Poll with all relations
+ */
+type PollTypeWithRelations = InferSelectModel<typeof Polls> & {
+    options: (InferSelectModel<typeof PollOptions> & {
+        votes: InferSelectModel<typeof PollVotes>[];
+    })[];
+    votes: InferSelectModel<typeof PollVotes>[];
+};
+
+/**
+ * Database class for managing polls
+ */
+export class Poll extends BaseInterface<typeof Polls, PollTypeWithRelations> {
+    public static $type: PollTypeWithRelations;
+
+    /**
+     * Reload the poll data from the database
+     */
+    public async reload(): Promise<void> {
+        const reloaded = await Poll.fromId(this.data.id);
+
+        if (!reloaded) {
+            throw new Error("Failed to reload poll");
+        }
+
+        this.data = reloaded.data;
+    }
+
+    /**
+     * Get a poll by ID
+     * @param id - The poll ID
+     * @returns The poll instance or null if not found
+     */
+    public static async fromId(id: string | null): Promise<Poll | null> {
+        if (!id) {
+            return null;
+        }
+
+        return await Poll.fromSql(eq(Polls.id, id));
+    }
+
+    /**
+     * Get a poll by note ID
+     * @param noteId - The note ID
+     * @returns The poll instance or null if not found
+     */
+    public static async fromNoteId(noteId: string): Promise<Poll | null> {
+        return await Poll.fromSql(eq(Polls.noteId, noteId));
+    }
+
+    /**
+     * Get multiple polls by IDs
+     * @param ids - Array of poll IDs
+     * @returns Array of poll instances
+     */
+    public static async fromIds(ids: string[]): Promise<Poll[]> {
+        return await Poll.manyFromSql(inArray(Polls.id, ids));
+    }
+
+    /**
+     * Execute SQL query to get a single poll with relations
+     * @param sql - SQL condition
+     * @returns Poll instance or null
+     */
+    protected static async fromSql(sql: any): Promise<Poll | null> {
+        const result = await db
+            .select()
+            .from(Polls)
+            .leftJoin(PollOptions, eq(Polls.id, PollOptions.pollId))
+            .leftJoin(PollVotes, eq(PollOptions.id, PollVotes.optionId))
+            .where(sql);
+
+        if (result.length === 0) {
+            return null;
+        }
+
+        // Group the results to build the poll object with options
+        const pollData = result[0].Polls;
+        const optionsMap = new Map<string, any>();
+        const votesData: InferSelectModel<typeof PollVotes>[] = [];
+
+        for (const row of result) {
+            if (row.PollOptions) {
+                if (!optionsMap.has(row.PollOptions.id)) {
+                    optionsMap.set(row.PollOptions.id, {
+                        ...row.PollOptions,
+                        votes: [],
+                    });
+                }
+
+                if (row.PollVotes) {
+                    optionsMap.get(row.PollOptions.id)!.votes.push(row.PollVotes);
+                    votesData.push(row.PollVotes);
+                }
+            }
+        }
+
+        const options = Array.from(optionsMap.values()).sort((a, b) => a.index - b.index);
+
+        const pollWithRelations: PollTypeWithRelations = {
+            ...pollData,
+            options,
+            votes: votesData,
+        };
+
+        return new Poll(pollWithRelations);
+    }
+
+    /**
+     * Execute SQL query to get multiple polls with relations
+     * @param sql - SQL condition
+     * @returns Array of poll instances
+     */
+    protected static async manyFromSql(sql: any): Promise<Poll[]> {
+        const result = await db
+            .select()
+            .from(Polls)
+            .leftJoin(PollOptions, eq(Polls.id, PollOptions.pollId))
+            .leftJoin(PollVotes, eq(PollOptions.id, PollVotes.optionId))
+            .where(sql);
+
+        if (result.length === 0) {
+            return [];
+        }
+
+        // Group by poll ID
+        const pollsMap = new Map<string, any>();
+
+        for (const row of result) {
+            const pollId = row.Polls.id;
+
+            if (!pollsMap.has(pollId)) {
+                pollsMap.set(pollId, {
+                    ...row.Polls,
+                    options: new Map(),
+                    votes: [],
+                });
+            }
+
+            const poll = pollsMap.get(pollId);
+
+            if (row.PollOptions) {
+                if (!poll.options.has(row.PollOptions.id)) {
+                    poll.options.set(row.PollOptions.id, {
+                        ...row.PollOptions,
+                        votes: [],
+                    });
+                }
+
+                if (row.PollVotes) {
+                    poll.options.get(row.PollOptions.id)!.votes.push(row.PollVotes);
+                    poll.votes.push(row.PollVotes);
+                }
+            }
+        }
+
+        return Array.from(pollsMap.values()).map((pollData) => {
+            const options = Array.from(pollData.options.values()).sort(
+                (a, b) => a.index - b.index,
+            );
+
+            return new Poll({
+                ...pollData,
+                options,
+                votes: pollData.votes,
+            });
+        });
+    }
+
+    /**
+     * Insert a new poll into the database
+     * @param pollData - Poll data to insert
+     * @param options - Poll options to insert
+     * @returns The inserted poll instance
+     */
+    public static async insert(
+        pollData: InferInsertModel<typeof Polls>,
+        options: string[],
+    ): Promise<Poll> {
+        return await db.transaction(async (tx) => {
+            // Insert the poll
+            const insertedPoll = (await tx.insert(Polls).values(pollData).returning())[0];
+
+            // Insert poll options
+            const optionInserts = options.map((title, index) => ({
+                id: crypto.randomUUID(),
+                pollId: insertedPoll.id,
+                title,
+                index,
+                votesCount: 0,
+            }));
+
+            await tx.insert(PollOptions).values(optionInserts);
+
+            // Return the poll with relations
+            const poll = await Poll.fromId(insertedPoll.id);
+            if (!poll) {
+                throw new Error("Failed to retrieve inserted poll");
+            }
+
+            return poll;
+        });
+    }
+
+    /**
+     * Check if the poll has expired
+     * @returns True if the poll has expired
+     */
+    public isExpired(): boolean {
+        if (!this.data.expiresAt) {
+            return false;
+        }
+
+        return new Date(this.data.expiresAt) < new Date();
+    }
+
+    /**
+     * Check if a user has voted in this poll
+     * @param userId - The user ID to check
+     * @returns True if the user has voted
+     */
+    public hasUserVoted(userId: string): boolean {
+        return this.data.votes.some((vote) => vote.userId === userId);
+    }
+
+    /**
+     * Get the vote options for a specific user
+     * @param userId - The user ID
+     * @returns Array of option indices the user voted for
+     */
+    public getUserVotes(userId: string): number[] {
+        const userVotes = this.data.votes.filter((vote) => vote.userId === userId);
+        return userVotes.map((vote) => {
+            const option = this.data.options.find((opt) => opt.id === vote.optionId);
+            return option?.index ?? -1;
+        }).filter((index) => index !== -1);
+    }
+
+    /**
+     * Convert poll to Mastodon API format
+     * @param userFetching - The user fetching the poll (to check if they voted)
+     * @returns Poll in Mastodon API format
+     */
+    public toApi(userFetching?: { id: string } | null): z.infer<typeof PollSchema> {
+        const voted = userFetching ? this.hasUserVoted(userFetching.id) : undefined;
+        const ownVotes = userFetching ? this.getUserVotes(userFetching.id) : undefined;
+
+        return {
+            id: this.data.id,
+            expires_at: this.data.expiresAt,
+            expired: this.isExpired(),
+            multiple: this.data.multiple,
+            votes_count: this.data.votesCount,
+            voters_count: this.data.votersCount,
+            options: this.data.options.map((option) => ({
+                title: option.title,
+                votes_count: this.data.hideTotals && !this.isExpired() ? null : option.votesCount,
+            })),
+            emojis: [], // TODO: Parse emojis from poll options
+            voted,
+            own_votes: ownVotes,
+        };
+    }
+}

classes/functions/status.ts

@@ -115,6 +115,16 @@ export const findManyNotes = async (
                             ...userRelations,
                         },
                     },
+                    poll: {
+                        with: {
+                            options: {
+                                with: {
+                                    votes: true,
+                                },
+                            },
+                            votes: true,
+                        },
+                    },
                 },
                 extras: {
                     pinned: userId
@@ -141,6 +151,16 @@ export const findManyNotes = async (
             },
             reply: true,
             quote: true,
+            poll: {
+                with: {
+                    options: {
+                        with: {
+                            votes: true,
+                        },
+                    },
+                    votes: true,
+                },
+            },
         },
         extras: {
             pinned: userId
@@ -176,6 +196,10 @@ export const findManyNotes = async (
         })),
         attachments: post.attachments.map((attachment) => attachment.media),
         emojis: (post.emojis ?? []).map((emoji) => emoji.emoji),
+        poll: post.poll ? {
+            ...post.poll,
+            options: post.poll.options.sort((a, b) => a.index - b.index),
+        } : null,
         reblog: post.reblog && {
             ...post.reblog,
             author: transformOutputToUserWithRelations(post.reblog.author),
@@ -187,6 +211,10 @@ export const findManyNotes = async (
                 (attachment) => attachment.media,
             ),
             emojis: (post.reblog.emojis ?? []).map((emoji) => emoji.emoji),
+            poll: post.reblog.poll ? {
+                ...post.reblog.poll,
+                options: post.reblog.poll.options.sort((a, b) => a.index - b.index),
+            } : null,
             pinned: Boolean(post.reblog.pinned),
             reblogged: Boolean(post.reblog.reblogged),
             muted: Boolean(post.reblog.muted),