🚧 wip: pglite instance

package.json

@@ -32,7 +32,8 @@
     "build-sitemap": "tsx ./scripts/buildSitemapIndex/index.ts",
     "build:analyze": "ANALYZE=true next build",
     "build:docker": "DOCKER=true next build && npm run build-sitemap",
-    "db:generate": "drizzle-kit generate",
+    "db:generate": "drizzle-kit generate && npm run db:generate-client",
+    "db:generate-client": "bun ./scripts/migrateClientDB/compile-migrations.ts",
     "db:migrate": "MIGRATION_DB=1 tsx ./scripts/migrateServerDB/index.ts",
     "db:push": "drizzle-kit push",
     "db:push-test": "NODE_ENV=test drizzle-kit push",
@@ -117,6 +118,7 @@
     "@clerk/themes": "^2.1.37",
     "@codesandbox/sandpack-react": "^2.19.9",
     "@cyntler/react-doc-viewer": "^1.17.0",
+    "@electric-sql/pglite": "^0.2.14",
     "@google/generative-ai": "^0.21.0",
     "@huggingface/inference": "^2.8.1",
     "@icons-pack/react-simple-icons": "9.6.0",

scripts/migrateClientDB/compile-migrations.ts

@@ -0,0 +1,11 @@
+import { readMigrationFiles } from 'drizzle-orm/migrator';
+import { join } from 'node:path';
+
+const dbBase = join(__dirname, '../../src/database');
+const migrationsFolder = join(dbBase, './server/migrations');
+
+const migrations = readMigrationFiles({ migrationsFolder: migrationsFolder });
+
+await Bun.write(join(dbBase, './client/migrations.json'), JSON.stringify(migrations));
+
+console.log('🏁Migrations compiled!');

src/database/client/db.ts

@@ -0,0 +1,13 @@
+import { IdbFs, PGlite } from '@electric-sql/pglite';
+import { vector } from '@electric-sql/pglite/vector';
+import { drizzle } from 'drizzle-orm/pglite';
+
+import * as schema from '../server/schemas/lobechat';
+
+const client = new PGlite({
+  extensions: { vector },
+  fs: new IdbFs('lobechat'),
+  relaxedDurability: true,
+});
+
+export const clientDB = drizzle({ client, schema });

src/database/client/migrate.ts

@@ -0,0 +1,14 @@
+import type { MigrationConfig } from 'drizzle-orm/migrator';
+
+import { clientDB } from './db';
+import migrations from './migrations.json';
+
+export const migrate = async () => {
+  // refs: https://github.com/drizzle-team/drizzle-orm/discussions/2532
+  // @ts-ignore
+  clientDB.dialect.migrate(migrations, clientDB.session, {
+    migrationsTable: 'drizzle_migrations',
+  } satisfies Omit<MigrationConfig, 'migrationsFolder'>);
+
+  return clientDB;
+};

src/database/server/models/__tests__/topic.test.ts

@@ -16,7 +16,7 @@ vi.mock('@/database/server/core/db', async () => ({
 
 const userId = 'topic-user-test';
 const sessionId = 'topic-session';
-const topicModel = new TopicModel(userId);
+const topicModel = new TopicModel(serverDB, userId);
 
 describe('TopicModel', () => {
   beforeEach(async () => {

src/database/server/models/topic.ts

@@ -1,7 +1,7 @@
 import { Column, count, inArray, sql } from 'drizzle-orm';
 import { and, desc, eq, exists, isNull, like, or } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server/core/db';
+import { LobeChatDatabase } from '@/database/server/type';
 
 import { NewMessage, TopicItem, messages, topics } from '../schemas/lobechat';
 import { idGenerator } from '../utils/idGenerator';
@@ -21,17 +21,19 @@ interface QueryTopicParams {
 
 export class TopicModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
   // **************** Query *************** //
 
   async query({ current = 0, pageSize = 9999, sessionId }: QueryTopicParams = {}) {
     const offset = current * pageSize;
 
     return (
-      serverDB
+      this.db
         .select({
           createdAt: topics.createdAt,
           favorite: topics.favorite,
@@ -52,13 +54,13 @@ export class TopicModel {
   }
 
   async findById(id: string) {
-    return serverDB.query.topics.findFirst({
+    return this.db.query.topics.findFirst({
       where: and(eq(topics.id, id), eq(topics.userId, this.userId)),
     });
   }
 
   async queryAll(): Promise<TopicItem[]> {
-    return serverDB
+    return this.db
       .select()
       .from(topics)
       .orderBy(topics.updatedAt)
@@ -74,31 +76,26 @@ export class TopicModel {
     const matchKeyword = (field: any) =>
       like(sql`lower(${field})` as unknown as Column, `%${keywordLowerCase}%`);
 
-    return serverDB.query.topics.findMany({
+    return this.db.query.topics.findMany({
       orderBy: [desc(topics.updatedAt)],
       where: and(
         eq(topics.userId, this.userId),
         this.matchSession(sessionId),
         or(
           matchKeyword(topics.title),
           exists(
-            serverDB
+            this.db
               .select()
               .from(messages)
-              .where(
-                and(
-                  eq(messages.topicId, topics.id),
-                  matchKeyword(messages.content)
-                )
-              ),
+              .where(and(eq(messages.topicId, topics.id), matchKeyword(messages.content))),
           ),
         ),
       ),
     });
   }
 
   async count() {
-    const result = await serverDB
+    const result = await this.db
       .select({
         count: count(),
       })
@@ -115,7 +112,7 @@ export class TopicModel {
     { messages: messageIds, ...params }: CreateTopicParams,
     id: string = this.genId(),
   ): Promise<TopicItem> {
-    return serverDB.transaction(async (tx) => {
+    return this.db.transaction(async (tx) => {
       // 在 topics 表中插入新的 topic
       const [topic] = await tx
         .insert(topics)
@@ -140,7 +137,7 @@ export class TopicModel {
 
   async batchCreate(topicParams: (CreateTopicParams & { id?: string })[]) {
     // 开始一个事务
-    return serverDB.transaction(async (tx) => {
+    return this.db.transaction(async (tx) => {
       // 在 topics 表中批量插入新的 topics
       const createdTopics = await tx
         .insert(topics)
@@ -173,7 +170,7 @@ export class TopicModel {
   }
 
   async duplicate(topicId: string, newTitle?: string) {
-    return serverDB.transaction(async (tx) => {
+    return this.db.transaction(async (tx) => {
       // find original topic
       const originalTopic = await tx.query.topics.findFirst({
         where: and(eq(topics.id, topicId), eq(topics.userId, this.userId)),
@@ -228,14 +225,14 @@ export class TopicModel {
    * Delete a session, also delete all messages and topics associated with it.
    */
   async delete(id: string) {
-    return serverDB.delete(topics).where(and(eq(topics.id, id), eq(topics.userId, this.userId)));
+    return this.db.delete(topics).where(and(eq(topics.id, id), eq(topics.userId, this.userId)));
   }
 
   /**
    * Deletes multiple topics based on the sessionId.
    */
   async batchDeleteBySessionId(sessionId?: string | null) {
-    return serverDB
+    return this.db
       .delete(topics)
       .where(and(this.matchSession(sessionId), eq(topics.userId, this.userId)));
   }
@@ -244,19 +241,19 @@ export class TopicModel {
    * Deletes multiple topics and all messages associated with them in a transaction.
    */
   async batchDelete(ids: string[]) {
-    return serverDB
+    return this.db
       .delete(topics)
       .where(and(inArray(topics.id, ids), eq(topics.userId, this.userId)));
   }
 
   async deleteAll() {
-    return serverDB.delete(topics).where(eq(topics.userId, this.userId));
+    return this.db.delete(topics).where(eq(topics.userId, this.userId));
   }
 
   // **************** Update *************** //
 
   async update(id: string, data: Partial<TopicItem>) {
-    return serverDB
+    return this.db
       .update(topics)
       .set({ ...data, updatedAt: new Date() })
       .where(and(eq(topics.id, id), eq(topics.userId, this.userId)))

src/database/server/schemas/lobechat/topic.ts

@@ -2,6 +2,8 @@
 import { boolean, jsonb, pgTable, text, unique } from 'drizzle-orm/pg-core';
 import { createInsertSchema } from 'drizzle-zod';
 
+import { ChatTopicMetadata } from '@/types/topic';
+
 import { idGenerator } from '../../utils/idGenerator';
 import { timestamps, timestamptz } from './_helpers';
 import { sessions } from './session';
@@ -21,7 +23,7 @@ export const topics = pgTable(
       .notNull(),
     clientId: text('client_id'),
     historySummary: text('history_summary'),
-    metadata: jsonb('metadata'),
+    metadata: jsonb('metadata').$type<ChatTopicMetadata | undefined>(),
     ...timestamps,
   },
   (t) => ({

src/database/server/type.ts

@@ -0,0 +1,7 @@
+import { PgliteDatabase } from 'drizzle-orm/pglite';
+
+import * as schema from '../server/schemas/lobechat';
+
+export type LobeChatDatabaseSchema = typeof schema;
+
+export type LobeChatDatabase = PgliteDatabase<LobeChatDatabaseSchema>;

src/layout/GlobalProvider/StoreInitialization.tsx

@@ -6,6 +6,7 @@ import { useTranslation } from 'react-i18next';
 import { createStoreUpdater } from 'zustand-utils';
 
 import { LOBE_URL_IMPORT_NAME } from '@/const/url';
+import { migrate } from '@/database/client/migrate';
 import { useIsMobile } from '@/hooks/useIsMobile';
 import { useEnabledDataSync } from '@/hooks/useSyncData';
 import { useAgentStore } from '@/store/agent';
@@ -90,6 +91,11 @@ const StoreInitialization = memo(() => {
     }
   }, [router, mobile]);
 
+  useEffect(() => {
+    migrate().then(() => {
+      console.log('migrate success!');
+    });
+  }, []);
   return null;
 });
 

src/server/routers/lambda/topic.ts

@@ -1,5 +1,6 @@
 import { z } from 'zod';
 
+import { serverDB } from '@/database/server';
 import { TopicModel } from '@/database/server/models/topic';
 import { authedProcedure, publicProcedure, router } from '@/libs/trpc';
 import { BatchTaskResult } from '@/types/service';
@@ -8,7 +9,7 @@ const topicProcedure = authedProcedure.use(async (opts) => {
   const { ctx } = opts;
 
   return opts.next({
-    ctx: { topicModel: new TopicModel(ctx.userId) },
+    ctx: { topicModel: new TopicModel(serverDB, ctx.userId) },
   });
 });
 
@@ -78,6 +79,7 @@ export const topicRouter = router({
     return ctx.topicModel.queryAll();
   }),
 
+  // TODO: this procedure should be used with authedProcedure
   getTopics: publicProcedure
     .input(
       z.object({
@@ -89,7 +91,7 @@ export const topicRouter = router({
     .query(async ({ input, ctx }) => {
       if (!ctx.userId) return [];
 
-      const topicModel = new TopicModel(ctx.userId);
+      const topicModel = new TopicModel(serverDB, ctx.userId);
 
       return topicModel.query(input);
     }),

src/services/topic/client.ts

@@ -1,11 +1,21 @@
-import { TopicModel } from '@/database/_deprecated/models/topic';
+import { clientDB } from '@/database/client/db';
+import { TopicModel } from '@/database/server/models/topic';
+import { useUserStore } from '@/store/user';
+import { userProfileSelectors } from '@/store/user/selectors';
 import { ChatTopic } from '@/types/topic';
 
 import { CreateTopicParams, ITopicService, QueryTopicParams } from './type';
 
 export class ClientService implements ITopicService {
+  private topicModel: TopicModel;
+  constructor() {
+    const userId = userProfileSelectors.userId(useUserStore.getState())!;
+
+    this.topicModel = new TopicModel(clientDB, userId);
+  }
+
   async createTopic(params: CreateTopicParams): Promise<string> {
-    const item = await TopicModel.create(params as any);
+    const item = await this.topicModel.create(params as any);
 
     if (!item) {
       throw new Error('topic create Error');
@@ -15,56 +25,52 @@ export class ClientService implements ITopicService {
   }
 
   async batchCreateTopics(importTopics: ChatTopic[]) {
-    return TopicModel.batchCreate(importTopics as any);
+    const data = await this.topicModel.batchCreate(importTopics as any);
+
+    return { added: data.length, ids: [], skips: [], success: true };
   }
 
   async cloneTopic(id: string, newTitle?: string) {
-    return TopicModel.duplicateTopic(id, newTitle);
+    const data = await this.topicModel.duplicate(id, newTitle);
+    return data.topic.id;
   }
 
   async getTopics(params: QueryTopicParams): Promise<ChatTopic[]> {
-    return TopicModel.query(params);
+    console.log('get',params);
+    const data = await this.topicModel.query(params);
+    console.log('topic:', data);
+    return data;
   }
 
   async searchTopics(keyword: string, sessionId?: string) {
-    return TopicModel.queryByKeyword(keyword, sessionId);
+    return this.topicModel.queryByKeyword(keyword, sessionId);
   }
 
   async getAllTopics() {
-    return TopicModel.queryAll();
+    return this.topicModel.queryAll();
   }
 
   async countTopics() {
-    return TopicModel.count();
-  }
-
-  async updateTopicFavorite(id: string, favorite?: boolean) {
-    return this.updateTopic(id, { favorite });
-  }
-
-  async updateTopicTitle(id: string, text: string) {
-    return this.updateTopic(id, { title: text });
+    return this.topicModel.count();
   }
 
   async updateTopic(id: string, data: Partial<ChatTopic>) {
-    const favorite = typeof data.favorite !== 'undefined' ? (data.favorite ? 1 : 0) : undefined;
-
-    return TopicModel.update(id, { ...data, favorite });
+    return this.topicModel.update(id, data);
   }
 
   async removeTopic(id: string) {
-    return TopicModel.delete(id);
+    return this.topicModel.delete(id);
   }
 
   async removeTopics(sessionId: string) {
-    return TopicModel.batchDeleteBySessionId(sessionId);
+    return this.topicModel.batchDeleteBySessionId(sessionId);
   }
 
   async batchRemoveTopics(topics: string[]) {
-    return TopicModel.batchDelete(topics);
+    return this.topicModel.batchDelete(topics);
   }
 
   async removeAllTopic() {
-    return TopicModel.clearTable();
+    return this.topicModel.deleteAll();
   }
 }