import pLimit from "p-limit"; import type { ZBookmarkList } from "@karakeep/shared/types/lists"; import type { ZTagBasic } from "@karakeep/shared/types/tags"; import { BookmarkTypes } from "@karakeep/shared/types/bookmarks"; import { logInfo, logStep, logSuccess } from "./log"; import { getTrpcClient, TrpcClient } from "./trpc"; import { waitUntil } from "./utils"; export interface SeedConfig { bookmarkCount: number; tagCount: number; listCount: number; concurrency: number; userCount: number; } export interface SeededBookmark { id: string; tags: ZTagBasic[]; listId?: string; title: string | null | undefined; } export interface UserSeedData { apiKey: string; trpc: TrpcClient; email: string; tags: ZTagBasic[]; lists: ZBookmarkList[]; bookmarks: SeededBookmark[]; } export interface SeedResult { users: UserSeedData[]; searchTerm: string; // For backwards compatibility, expose the first user's data apiKey: string; trpc: TrpcClient; tags: ZTagBasic[]; lists: ZBookmarkList[]; bookmarks: SeededBookmark[]; } const TOPICS = [ "performance", "search", "reading", "workflow", "api", "workers", "backend", "frontend", "productivity", "cli", ]; async function seedUserData( authlessClient: TrpcClient, userIndex: number, config: SeedConfig, timestamp: number, ): Promise { const email = `benchmarks+${timestamp}+user${userIndex}@example.com`; const password = "benchmarks1234"; logStep(`Creating user ${userIndex + 1}/${config.userCount}`); await authlessClient.users.create.mutate({ name: `Benchmark User ${userIndex + 1}`, email, password, confirmPassword: password, }); const { key } = await authlessClient.apiKeys.exchange.mutate({ email, password, keyName: `benchmark-key-${userIndex}`, }); const trpc = getTrpcClient(key); logSuccess(`User ${userIndex + 1} ready`); logStep(`Creating ${config.tagCount} tags for user ${userIndex + 1}`); const tags: ZTagBasic[] = []; for (let i = 0; i < config.tagCount; i++) { const tag = await trpc.tags.create.mutate({ name: `user${userIndex}-topic-${i + 1}`, }); tags.push(tag); } logSuccess(`Tags created for user ${userIndex + 1}`); logStep(`Creating ${config.listCount} lists for user ${userIndex + 1}`); const lists: ZBookmarkList[] = []; for (let i = 0; i < config.listCount; i++) { const list = await trpc.lists.create.mutate({ name: `User ${userIndex + 1} List ${i + 1}`, description: `Auto-generated benchmark list #${i + 1} for user ${userIndex + 1}`, icon: "bookmark", }); lists.push(list); } logSuccess(`Lists created for user ${userIndex + 1}`); logStep( `Creating ${config.bookmarkCount} bookmarks for user ${userIndex + 1}`, ); const limit = pLimit(config.concurrency); const bookmarks: SeededBookmark[] = []; await Promise.all( Array.from({ length: config.bookmarkCount }).map((_, index) => limit(async () => { const topic = TOPICS[index % TOPICS.length]; const createdAt = new Date(Date.now() - index * 3000); const bookmark = await trpc.bookmarks.createBookmark.mutate({ type: BookmarkTypes.LINK, url: `https://example.com/user${userIndex}/${topic}/${index}`, title: `User ${userIndex + 1} ${topic} article ${index}`, source: "api", summary: `Benchmark dataset entry about ${topic} for user ${userIndex + 1}.`, favourited: index % 7 === 0, archived: index % 11 === 0, createdAt, }); const primaryTag = tags[index % tags.length]; const secondaryTag = tags[(index + 5) % tags.length]; const attachedTags = [primaryTag, secondaryTag]; await trpc.bookmarks.updateTags.mutate({ bookmarkId: bookmark.id, attach: attachedTags.map((tag) => ({ tagId: tag.id, tagName: tag.name, })), detach: [], }); let listId: string | undefined; if (lists.length > 0) { const list = lists[index % lists.length]; await trpc.lists.addToList.mutate({ listId: list.id, bookmarkId: bookmark.id, }); listId = list.id; } bookmarks.push({ id: bookmark.id, tags: attachedTags, listId, title: bookmark.title, }); }), ), ); logSuccess(`Bookmarks created for user ${userIndex + 1}`); return { apiKey: key, trpc, email, tags, lists, bookmarks, }; } export async function seedData(config: SeedConfig): Promise { const authlessClient = getTrpcClient(); const timestamp = Date.now(); logInfo(`Seeding data for ${config.userCount} users`); const users: UserSeedData[] = []; // Create all users sequentially to avoid race conditions for (let i = 0; i < config.userCount; i++) { const userData = await seedUserData(authlessClient, i, config, timestamp); users.push(userData); } const searchTerm = "benchmark"; logStep("Waiting for search index to be ready"); // Use the first user's client to check search readiness await waitUntil( async () => { const results = await users[0].trpc.bookmarks.searchBookmarks.query({ text: searchTerm, limit: 1, }); return results.bookmarks.length > 0; }, "search data to be indexed", 120_000, 2_000, ); logSuccess("Search index warmed up"); const totalBookmarks = users.reduce( (sum, user) => sum + user.bookmarks.length, 0, ); const totalTags = users.reduce((sum, user) => sum + user.tags.length, 0); const totalLists = users.reduce((sum, user) => sum + user.lists.length, 0); logInfo( `Seeded ${totalBookmarks} bookmarks across ${totalTags} tags and ${totalLists} lists for ${config.userCount} users`, ); // Return first user's data for backwards compatibility const firstUser = users[0]; return { users, searchTerm, apiKey: firstUser.apiKey, trpc: firstUser.trpc, tags: firstUser.tags, lists: firstUser.lists, bookmarks: firstUser.bookmarks, }; }