✨ feat: refactor client db to pglite

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": "tsx ./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,14 @@
+import { readMigrationFiles } from 'drizzle-orm/migrator';
+import { writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+
+const dbBase = join(__dirname, '../../src/database');
+const migrationsFolder = join(dbBase, './migrations');
+const migrations = readMigrationFiles({ migrationsFolder: migrationsFolder });
+
+writeFileSync(
+  join(dbBase, './client/migrations.json'),
+  JSON.stringify(migrations, null, 2), // null, 2 adds indentation for better readability
+);
+
+console.log('🏁 client migrations.json compiled!');

src/database/client/db.ts

@@ -0,0 +1,33 @@
+import { IdbFs, MemoryFS, PGlite } from '@electric-sql/pglite';
+import { vector } from '@electric-sql/pglite/vector';
+import { PgliteDatabase, drizzle } from 'drizzle-orm/pglite';
+
+import * as schema from '../schemas';
+
+let dbInstance: ReturnType<typeof drizzle>;
+
+export function getClientDB() {
+  // 如果已经初始化过，直接返回实例
+  if (dbInstance) return dbInstance;
+
+  // 服务端环境
+  if (typeof window === 'undefined') {
+    const db = new PGlite({
+      extensions: { vector },
+      fs: new MemoryFS('lobechat'),
+    });
+    return drizzle({ client: db, schema });
+  }
+
+  // 客户端环境
+  const db = new PGlite({
+    extensions: { vector },
+    fs: new IdbFs('lobechat'),
+    relaxedDurability: true,
+  });
+
+  dbInstance = drizzle({ client: db, schema });
+  return dbInstance;
+}
+
+export const clientDB = getClientDB() as unknown as PgliteDatabase<typeof schema>;

src/database/client/migrate.ts

@@ -0,0 +1,24 @@
+import { clientDB } from './db';
+import migrations from './migrations.json';
+
+export const migrate = async () => {
+  //prevent multiple schema migrations to be run
+  let isLocalDBSchemaSynced = false;
+
+  if (!isLocalDBSchemaSynced) {
+    const start = Date.now();
+    try {
+      // refs: https://github.com/drizzle-team/drizzle-orm/discussions/2532
+      // @ts-ignore
+      await clientDB.dialect.migrate(migrations, clientDB.session, {});
+      isLocalDBSchemaSynced = true;
+
+      console.info(`✅ Local database ready in ${Date.now() - start}ms`);
+    } catch (cause) {
+      console.error('❌ Local database schema migration failed', cause);
+      throw cause;
+    }
+  }
+
+  return clientDB;
+};

src/database/server/models/session.ts

@@ -229,6 +229,8 @@ export class SessionModel {
   }
 
   async updateConfig(id: string, data: Partial<AgentItem>) {
+    if (Object.keys(data).length === 0) return;
+
     return this.db
       .update(agents)
       .set(data)

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/file.ts

@@ -32,9 +32,7 @@ export const fileRouter = router({
     }),
 
   createFile: fileProcedure
-    .input(
-      UploadFileSchema.omit({ data: true, saveMode: true, url: true }).extend({ url: z.string() }),
-    )
+    .input(UploadFileSchema.omit({ url: true }).extend({ url: z.string() }))
     .mutation(async ({ ctx, input }) => {
       const { isExist } = await ctx.fileModel.checkHash(input.hash!);
 

src/services/baseClientService/index.ts

@@ -0,0 +1,19 @@
+const getClientDBUserId = () => {
+  if (typeof window === 'undefined') return undefined;
+
+  return window.__lobeClientUserId;
+};
+
+const FALLBACK_CLIENT_DB_USER_ID = 'DEFAULT_LOBE_CHAT_USER';
+
+export class BaseClientService {
+  private readonly fallbackUserId: string;
+
+  protected get userId(): string {
+    return getClientDBUserId() || this.fallbackUserId;
+  }
+
+  constructor(userId?: string) {
+    this.fallbackUserId = userId || FALLBACK_CLIENT_DB_USER_ID;
+  }
+}

src/services/file/ClientS3/index.ts

@@ -0,0 +1,56 @@
+import { createStore, del, get, set } from 'idb-keyval';
+
+const BROWSER_S3_DB_NAME = 'lobechat-local-s3';
+
+class BrowserS3Storage {
+  private store;
+
+  constructor() {
+    // skip server-side rendering
+    if (typeof window === 'undefined') return;
+
+    this.store = createStore(BROWSER_S3_DB_NAME, 'objects');
+  }
+
+  /**
+   * 上传文件
+   * @param key 文件 hash
+   * @param file File 对象
+   */
+  async putObject(key: string, file: File): Promise<void> {
+    try {
+      const data = await file.arrayBuffer();
+      await set(key, { data, name: file.name, type: file.type }, this.store);
+    } catch (e) {
+      throw new Error(`Failed to put file ${file.name}: ${(e as Error).message}`);
+    }
+  }
+
+  /**
+   * 获取文件
+   * @param key 文件 hash
+   * @returns File 对象
+   */
+  async getObject(key: string): Promise<File | undefined> {
+    try {
+      const res = await get<{ data: ArrayBuffer; name: string; type: string }>(key, this.store);
+      return new File([res!.data], res!.name, { type: res?.type });
+    } catch (e) {
+      throw new Error(`Failed to get object (key=${key}): ${(e as Error).message}`);
+    }
+  }
+
+  /**
+   * 删除文件
+   * @param key 文件 hash
+   */
+  async deleteObject(key: string): Promise<void> {
+    try {
+      await del(key, this.store);
+    } catch (e) {
+      throw new Error(`Failed to delete object (key=${key}): ${(e as Error).message}`);
+    }
+  }
+}
+
+export const clientS3Storage = new BrowserS3Storage();

src/services/file/client.test.ts

@@ -1,110 +1,109 @@
-import { Mock, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import { eq } from 'drizzle-orm';
+import { Mock, beforeEach, describe, expect, it } from 'vitest';
 
-import { fileEnv } from '@/config/file';
-import { FileModel } from '@/database/_deprecated/models/file';
-import { DB_File } from '@/database/_deprecated/schemas/files';
-import { serverConfigSelectors } from '@/store/serverConfig/selectors';
-import { createServerConfigStore } from '@/store/serverConfig/store';
+import { clientDB } from '@/database/client/db';
+import { migrate } from '@/database/client/migrate';
+import { files, globalFiles, users } from '@/database/schemas';
+import { clientS3Storage } from '@/services/file/ClientS3';
+import { UploadFileParams } from '@/types/files';
 
 import { ClientService } from './client';
 
-const fileService = new ClientService();
+const userId = 'file-user';
 
-beforeAll(() => {
-  createServerConfigStore();
-});
-// Mocks for the FileModel
-vi.mock('@/database/_deprecated/models/file', () => ({
-  FileModel: {
-    create: vi.fn(),
-    delete: vi.fn(),
-    findById: vi.fn(),
-    clear: vi.fn(),
-  },
-}));
-
-let s3Domain: string;
-
-vi.mock('@/config/file', () => ({
-  fileEnv: {
-    get NEXT_PUBLIC_S3_DOMAIN() {
-      return s3Domain;
-    },
-  },
-}));
-
-// Mocks for the URL and Blob objects
-global.URL.createObjectURL = vi.fn();
-global.Blob = vi.fn();
-
-beforeEach(() => {
-  // Reset all mocks before each test
-  vi.resetAllMocks();
-  s3Domain = '';
+const fileService = new ClientService(userId);
+
+const mockFile = {
+  name: 'mock.png',
+  fileType: 'image/png',
+  size: 1,
+  url: '',
+};
+
+beforeEach(async () => {
+  await migrate();
+
+  await clientDB.delete(users);
+  // 创建测试数据
+  await clientDB.transaction(async (tx) => {
+    await tx.insert(users).values({ id: userId });
+  });
 });
 
 describe('FileService', () => {
   it('createFile should save the file to the database', async () => {
-    const localFile: DB_File = {
+    const localFile: UploadFileParams = {
       name: 'test',
-      data: new ArrayBuffer(1),
       fileType: 'image/png',
-      saveMode: 'local',
+      url: '',
       size: 1,
+      hash: '123',
     };
 
-    (FileModel.create as Mock).mockResolvedValue(localFile);
+    await clientS3Storage.putObject(
+      '123',
+      new File([new ArrayBuffer(1)], 'test.png', { type: 'image/png' }),
+    );
 
     const result = await fileService.createFile(localFile);
 
-    expect(FileModel.create).toHaveBeenCalledWith(localFile);
-    expect(result).toEqual({ url: 'data:image/png;base64,AA==' });
+    expect(result).toMatchObject({ url: 'data:image/png;base64,AA==' });
   });
 
   it('removeFile should delete the file from the database', async () => {
     const fileId = '1';
-    (FileModel.delete as Mock).mockResolvedValue(true);
+    await clientDB.insert(files).values({ id: fileId, userId, ...mockFile });
+
+    await fileService.removeFile(fileId);
 
-    const result = await fileService.removeFile(fileId);
+    const result = await clientDB.query.files.findFirst({
+      where: eq(files.id, fileId),
+    });
 
-    expect(FileModel.delete).toHaveBeenCalledWith(fileId);
-    expect(result).toBe(true);
+    expect(result).toBeUndefined();
   });
 
   describe('getFile', () => {
     it('should retrieve and convert local file info to FilePreview', async () => {
-      const fileId = '1';
-      const fileData = {
-        name: 'test',
-        data: new ArrayBuffer(1),
+      const fileId = 'rwlijweled';
+      const file = {
         fileType: 'image/png',
-        saveMode: 'local',
         size: 1,
-        createdAt: 1,
-        updatedAt: 2,
-      } as DB_File;
+        name: 'test.png',
+        url: 'idb://12312/abc.png',
+        hashId: '123tttt',
+      };
+
+      await clientDB.insert(globalFiles).values(file);
+
+      await clientDB.insert(files).values({
+        id: fileId,
+        userId,
+        ...file,
+        createdAt: new Date(1),
+        updatedAt: new Date(2),
+        fileHash: file.hashId,
+      });
 
-      (FileModel.findById as Mock).mockResolvedValue(fileData);
-      (global.URL.createObjectURL as Mock).mockReturnValue('blob:test');
-      (global.Blob as Mock).mockImplementation(() => ['test']);
+      await clientS3Storage.putObject(
+        file.hashId,
+        new File([new ArrayBuffer(1)], file.name, { type: file.fileType }),
+      );
 
       const result = await fileService.getFile(fileId);
 
-      expect(FileModel.findById).toHaveBeenCalledWith(fileId);
-      expect(result).toEqual({
+      expect(result).toMatchObject({
         createdAt: new Date(1),
-        id: '1',
+        id: 'rwlijweled',
         size: 1,
         type: 'image/png',
-        name: 'test',
-        url: 'blob:test',
+        name: 'test.png',
         updatedAt: new Date(2),
       });
     });
 
     it('should throw an error when the file is not found', async () => {
       const fileId = 'non-existent';
-      (FileModel.findById as Mock).mockResolvedValue(null);
 
       const getFilePromise = fileService.getFile(fileId);