@lobehub/chat 1.74.1 → 1.74.3

package/src/database/client/db.ts

@@ -1,7 +1,14 @@
+import { sql } from 'drizzle-orm';
 import { PgliteDatabase, drizzle } from 'drizzle-orm/pglite';
 import { Md5 } from 'ts-md5';
 
-import { ClientDBLoadingProgress, DatabaseLoadingState } from '@/types/clientDB';
+import { DrizzleMigrationModel } from '@/database/models/drizzleMigration';
+import {
+  ClientDBLoadingProgress,
+  DatabaseLoadingState,
+  MigrationSQL,
+  MigrationTableItem,
+} from '@/types/clientDB';
 import { sleep } from '@/utils/sleep';
 
 import * as schema from '../schemas';
@@ -9,10 +16,17 @@ import migrations from './migrations.json';
 
 const pgliteSchemaHashCache = 'LOBE_CHAT_PGLITE_SCHEMA_HASH';
 
+const DB_NAME = 'lobechat';
 type DrizzleInstance = PgliteDatabase<typeof schema>;
 
+interface onErrorState {
+  error: Error;
+  migrationTableItems: MigrationTableItem[];
+  migrationsSQL: MigrationSQL[];
+}
+
 export interface DatabaseLoadingCallbacks {
-  onError?: (error: Error) => void;
+  onError?: (error: onErrorState) => void;
   onProgress?: (progress: ClientDBLoadingProgress) => void;
   onStateChange?: (state: DatabaseLoadingState) => void;
 }
@@ -152,8 +166,28 @@ export class DatabaseManager {
       hash = Md5.hashStr(JSON.stringify(migrations));
       // if hash is the same, no need to migrate
       if (hash === cacheHash) {
-        this.isLocalDBSchemaSynced = true;
-        return this.db;
+        try {
+          // 检查数据库中是否存在表
+          // 这里使用 pg_tables 系统表查询用户表数量
+          const tablesResult = await this.db.execute(
+            sql`
+              SELECT COUNT(*) as table_count
+              FROM information_schema.tables
+              WHERE table_schema = 'public'
+            `,
+          );
+
+          const tableCount = parseInt((tablesResult.rows[0] as any).table_count || '0', 10);
+
+          // 如果表数量大于0，则认为数据库已正确初始化
+          if (tableCount > 0) {
+            this.isLocalDBSchemaSynced = true;
+            return this.db;
+          }
+        } catch (error) {
+          console.warn('Error checking table existence, proceeding with migration', error);
+          // 如果查询失败，继续执行迁移以确保安全
+        }
       }
     }
 
@@ -204,13 +238,11 @@ export class DatabaseManager {
 
         let db: typeof PGlite;
 
-        const dbName = 'lobechat';
-
         // make db as web worker if worker is available
         // https://github.com/lobehub/lobe-chat/issues/5785
         if (typeof Worker !== 'undefined' && typeof navigator.locks !== 'undefined') {
           db = await initPgliteWorker({
-            dbName,
+            dbName: DB_NAME,
             fsBundle: fsBundle as Blob,
             vectorBundlePath: DatabaseManager.VECTOR_CDN_URL,
             wasmModule,
@@ -219,7 +251,7 @@ export class DatabaseManager {
           // in edge runtime or test runtime, we don't have worker
           db = new PGlite({
             extensions: { vector },
-            fs: typeof window === 'undefined' ? new MemoryFS(dbName) : new IdbFs(dbName),
+            fs: typeof window === 'undefined' ? new MemoryFS(DB_NAME) : new IdbFs(DB_NAME),
             relaxedDurability: true,
             wasmModule,
           });
@@ -241,10 +273,25 @@ export class DatabaseManager {
         this.initPromise = null;
         this.callbacks?.onStateChange?.(DatabaseLoadingState.Error);
         const error = e as Error;
+
+        // 查询迁移表数据
+        let migrationsTableData: MigrationTableItem[] = [];
+        try {
+          // 尝试查询迁移表
+          const drizzleMigration = new DrizzleMigrationModel(this.db as any);
+          migrationsTableData = await drizzleMigration.getMigrationList();
+        } catch (queryError) {
+          console.error('Failed to query migrations table:', queryError);
+        }
+
         this.callbacks?.onError?.({
-          message: error.message,
-          name: error.name,
-          stack: error.stack,
+          error: {
+            message: error.message,
+            name: error.name,
+            stack: error.stack,
+          },
+          migrationTableItems: migrationsTableData,
+          migrationsSQL: migrations,
         });
 
         console.error(error);
@@ -271,6 +318,38 @@ export class DatabaseManager {
       },
     });
   }
+
+  async resetDatabase(): Promise<void> {
+    // 删除 IndexedDB 数据库
+    return new Promise<void>((resolve, reject) => {
+      // 检查 IndexedDB 是否可用
+      if (typeof indexedDB === 'undefined') {
+        console.warn('IndexedDB is not available, cannot delete database');
+        resolve();
+        return;
+      }
+
+      const dbName = `/pglite/${DB_NAME}`;
+      const request = indexedDB.deleteDatabase(dbName);
+
+      request.onsuccess = () => {
+        console.log(`✅ Database '${dbName}' reset successfully`);
+
+        // 清除本地存储的模式哈希
+        if (typeof localStorage !== 'undefined') {
+          localStorage.removeItem(pgliteSchemaHashCache);
+        }
+
+        resolve();
+      };
+
+      // eslint-disable-next-line unicorn/prefer-add-event-listener
+      request.onerror = (event) => {
+        console.error('❌ Error resetting database:', event);
+        reject(new Error(`Failed to reset database '${dbName}'`));
+      };
+    });
+  }
 }
 
 // 导出单例
@@ -282,3 +361,15 @@ export const clientDB = dbManager.createProxy();
 // 导出初始化方法，供应用启动时使用
 export const initializeDB = (callbacks?: DatabaseLoadingCallbacks) =>
   dbManager.initialize(callbacks);
+
+export const resetClientDatabase = async () => {
+  await dbManager.resetDatabase();
+};
+
+export const updateMigrationRecord = async (migrationHash: string) => {
+  await clientDB.execute(
+    sql`INSERT INTO "drizzle"."__drizzle_migrations" ("hash", "created_at") VALUES (${migrationHash}, ${Date.now()});`,
+  );
+
+  await initializeDB();
+};

package/src/database/client/migrations.json

@@ -223,10 +223,7 @@
     "hash": "9646161fa041354714f823d726af27247bcd6e60fa3be5698c0d69f337a5700b"
   },
   {
-    "sql": [
-      "DROP TABLE \"user_budgets\";",
-      "\nDROP TABLE \"user_subscriptions\";"
-    ],
+    "sql": ["DROP TABLE \"user_budgets\";", "\nDROP TABLE \"user_subscriptions\";"],
     "bps": true,
     "folderMillis": 1729699958471,
     "hash": "7dad43a2a25d1aec82124a4e53f8d82f8505c3073f23606c1dc5d2a4598eacf9"
@@ -298,9 +295,7 @@
     "hash": "845a692ceabbfc3caf252a97d3e19a213bc0c433df2689900135f9cfded2cf49"
   },
   {
-    "sql": [
-      "ALTER TABLE \"messages\" ADD COLUMN \"reasoning\" jsonb;"
-    ],
+    "sql": ["ALTER TABLE \"messages\" ADD COLUMN \"reasoning\" jsonb;"],
     "bps": true,
     "folderMillis": 1737609172353,
     "hash": "2cb36ae4fcdd7b7064767e04bfbb36ae34518ff4bb1b39006f2dd394d1893868"
@@ -385,5 +380,40 @@
     "bps": true,
     "folderMillis": 1742269437903,
     "hash": "89e91285be422d5f44511c7405f57b57f8dfda4f0304126ae9b0f266e5fd60f1"
+  },
+  {
+    "sql": [
+      "ALTER TABLE \"session_groups\" DROP CONSTRAINT \"session_group_client_id_user_unique\";",
+      "\nALTER TABLE \"sessions\" DROP CONSTRAINT \"sessions_client_id_user_id_unique\";",
+      "\nALTER TABLE \"topics\" DROP CONSTRAINT \"topic_client_id_user_id_unique\";",
+      "\nALTER TABLE \"agents\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"files\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"knowledge_bases\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"message_plugins\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"message_queries\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"message_tts\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"message_translates\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"chunks\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"embeddings\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"unstructured_chunks\" ADD COLUMN \"client_id\" text;",
+      "\nALTER TABLE \"threads\" ADD COLUMN \"client_id\" text;",
+      "\nCREATE UNIQUE INDEX \"client_id_user_id_unique\" ON \"agents\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"files_client_id_user_id_unique\" ON \"files\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"knowledge_bases_client_id_user_id_unique\" ON \"knowledge_bases\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"message_plugins_client_id_user_id_unique\" ON \"message_plugins\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"message_queries_client_id_user_id_unique\" ON \"message_queries\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"message_tts_client_id_user_id_unique\" ON \"message_tts\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"message_translates_client_id_user_id_unique\" ON \"message_translates\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"chunks_client_id_user_id_unique\" ON \"chunks\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"embeddings_client_id_user_id_unique\" ON \"embeddings\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"unstructured_chunks_client_id_user_id_unique\" ON \"unstructured_chunks\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"session_groups_client_id_user_id_unique\" ON \"session_groups\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"sessions_client_id_user_id_unique\" ON \"sessions\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"threads_client_id_user_id_unique\" ON \"threads\" USING btree (\"client_id\",\"user_id\");",
+      "\nCREATE UNIQUE INDEX \"topics_client_id_user_id_unique\" ON \"topics\" USING btree (\"client_id\",\"user_id\");"
+    ],
+    "bps": true,
+    "folderMillis": 1742616026643,
+    "hash": "f885f1b1467b6564498f0d54ac2366c2d596f59bdd63a0feb63810c2ba346ae5"
   }
-]
+]

package/src/database/migrations/0018_add_client_id_for_entities.sql

@@ -0,0 +1,32 @@
+ALTER TABLE "session_groups" DROP CONSTRAINT IF EXISTS "session_group_client_id_user_unique";--> statement-breakpoint
+ALTER TABLE "sessions" DROP CONSTRAINT IF EXISTS "sessions_client_id_user_id_unique";--> statement-breakpoint
+ALTER TABLE "topics" DROP CONSTRAINT IF EXISTS "topic_client_id_user_id_unique";--> statement-breakpoint
+
+-- add client_id column
+ALTER TABLE "agents" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "files" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "knowledge_bases" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "message_plugins" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "message_queries" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "message_tts" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "message_translates" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "chunks" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "embeddings" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "unstructured_chunks" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+ALTER TABLE "threads" ADD COLUMN IF NOT EXISTS "client_id" text;--> statement-breakpoint
+
+-- Create unique index（using IF NOT EXISTS）
+CREATE UNIQUE INDEX IF NOT EXISTS "client_id_user_id_unique" ON "agents" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "files_client_id_user_id_unique" ON "files" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "knowledge_bases_client_id_user_id_unique" ON "knowledge_bases" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "message_plugins_client_id_user_id_unique" ON "message_plugins" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "message_queries_client_id_user_id_unique" ON "message_queries" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "message_tts_client_id_user_id_unique" ON "message_tts" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "message_translates_client_id_user_id_unique" ON "message_translates" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "chunks_client_id_user_id_unique" ON "chunks" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "embeddings_client_id_user_id_unique" ON "embeddings" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "unstructured_chunks_client_id_user_id_unique" ON "unstructured_chunks" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "session_groups_client_id_user_id_unique" ON "session_groups" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "sessions_client_id_user_id_unique" ON "sessions" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "threads_client_id_user_id_unique" ON "threads" USING btree ("client_id","user_id");--> statement-breakpoint
+CREATE UNIQUE INDEX IF NOT EXISTS "topics_client_id_user_id_unique" ON "topics" USING btree ("client_id","user_id");