🚧 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": "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,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 '../schemas';
+
+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,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/migrations/0006_add_knowledge_base.sql

@@ -121,7 +121,7 @@ CREATE TABLE IF NOT EXISTS "unstructured_chunks" (
 	"file_id" varchar
 );
 --> statement-breakpoint
-ALTER TABLE "files_to_messages" RENAME TO "messages_files";
+ALTER TABLE "files_to_messages" RENAME TO "messages_files";--> statement-breakpoint
 DROP TABLE "files_to_agents";--> statement-breakpoint
 ALTER TABLE "files" ADD COLUMN "file_hash" varchar(64);--> statement-breakpoint
 ALTER TABLE "files" ADD COLUMN "chunk_task_id" uuid;--> statement-breakpoint

src/database/migrations/0007_fix_embedding_table.sql

@@ -3,16 +3,20 @@ CREATE TEMP TABLE embeddings_temp AS
 SELECT DISTINCT ON (chunk_id) *
 FROM embeddings
 ORDER BY chunk_id, random();
+--> statement-breakpoint
 
 -- step 2: delete all rows from the original table
 DELETE FROM embeddings;
+--> statement-breakpoint
 
 -- step 3: insert the rows we want to keep back into the original table
 INSERT INTO embeddings
 SELECT * FROM embeddings_temp;
+--> statement-breakpoint
 
 -- step 4: drop the temporary table
 DROP TABLE embeddings_temp;
+--> statement-breakpoint
 
 -- step 5: now it's safe to add the unique constraint
 ALTER TABLE "embeddings" ADD CONSTRAINT "embeddings_chunk_id_unique" UNIQUE("chunk_id");

src/database/server/models/topic.ts

@@ -2,9 +2,9 @@ import { Column, count, inArray, sql } from 'drizzle-orm';
 import { and, desc, eq, exists, isNull, like, or } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
+import { idGenerator } from '@/database/utils/idGenerator';
 
 import { NewMessage, TopicItem, messages, topics } from '../../schemas';
-import { idGenerator } from '@/database/utils/idGenerator';
 
 export interface CreateTopicParams {
   favorite?: boolean;

src/database/server/models/user.ts

@@ -94,7 +94,7 @@ export class UserModel {
     return this.db.delete(userSettings).where(eq(userSettings.id, this.userId));
   }
 
-  async updateSetting(value: Partial<UserSettings>) {
+  async updateSetting(value: DeepPartial<UserSettings>) {
     const { keyVaults, ...res } = value;
 
     // Encrypt keyVaults

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/services/import/client.ts

@@ -32,7 +32,7 @@ export class ClientService {
     }
 
     if (sessionGroups.length > 0) {
-      const res = await SessionGroupModel.batchCreate(sessionGroups);
+      const res = await SessionGroupModel.batchCreate(sessionGroups as any);
       sessionGroupResult = this.mapImportResult(res);
     }