♻️ refactor: refactor the db to context inject mode (lobehub#7255)

* refactor with new db init mode

* fix tests

* fix tests

* move the separate index

* fix tests

* fix tests

* fix db issue

* fix db

* refactor to clean

* Update index.ts

* fix error

* fix the exist inbox slug session

* fix the tests

Author: arvinxx

src/database/core/db-adaptor.ts

@@ -1,6 +1,24 @@
-// import { isDesktop } from '@/const/version';
 import { getDBInstance } from '@/database/core/web-server';
+import { LobeChatDatabase } from '@/database/type';
 
-// import { getPgliteInstance } from './electron';
+/**
+ * 懒加载数据库实例
+ * 避免每次模块导入时都初始化数据库
+ */
+let cachedDB: LobeChatDatabase | null = null;
+
+export const getServerDB = async (): Promise<LobeChatDatabase> => {
+  // 如果已经有缓存的实例，直接返回
+  if (cachedDB) return cachedDB;
+
+  try {
+    // 根据环境选择合适的数据库实例
+    cachedDB = getDBInstance();
+    return cachedDB;
+  } catch (error) {
+    console.error('❌ Failed to initialize database:', error);
+    throw error;
+  }
+};
 
 export const serverDB = getDBInstance();

src/database/models/__tests__/aiProvider.test.ts

@@ -4,6 +4,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { LobeChatDatabase } from '@/database/type';
 import { ModelProvider } from '@/libs/agent-runtime';
+import { sleep } from '@/utils/sleep';
 
 import { aiProviders, users } from '../../schemas';
 import { AiProviderModel } from '../aiProvider';
@@ -96,6 +97,7 @@ describe('AiProviderModel', () => {
   describe('query', () => {
     it('should query ai providers for the user', async () => {
       await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
+      await sleep(10);
       await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
 
       const userGroups = await aiProviderModel.query();

src/database/models/session.ts

@@ -193,6 +193,14 @@ export class SessionModel {
     type: 'agent' | 'group';
   }): Promise<SessionItem> => {
     return this.db.transaction(async (trx) => {
+      if (slug) {
+        const existResult = await trx.query.sessions.findFirst({
+          where: and(eq(sessions.slug, slug), eq(sessions.userId, this.userId)),
+        });
+
+        if (existResult) return existResult;
+      }
+
       const newAgents = await trx
         .insert(agents)
         .values({

src/database/models/user.ts

@@ -85,7 +85,13 @@ export class UserModel {
     const state = result[0];
 
     // Decrypt keyVaults
-    const decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
+    let decryptKeyVaults = {};
+
+    try {
+      decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
+    } catch {
+      /* empty */
+    }
 
     const settings: DeepPartial<UserSettings> = {
       defaultAgent: state.settingsDefaultAgent || {},

src/database/server/index.ts

@@ -1 +1 @@
-export { serverDB } from '../core/db-adaptor';
+export { getServerDB, serverDB } from '../core/db-adaptor';

src/libs/trpc/async/index.ts

@@ -1,10 +1,20 @@
+import { getServerDB } from '@/database/core/db-adaptor';
+
 import { asyncAuth } from './asyncAuth';
 import { asyncTrpc } from './init';
 
 export const publicProcedure = asyncTrpc.procedure;
 
 export const asyncRouter = asyncTrpc.router;
 
-export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth);
+export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth).use(
+  asyncTrpc.middleware(async (opts) => {
+    const serverDB = await getServerDB();
+
+    return opts.next({
+      ctx: { serverDB },
+    });
+  }),
+);
 
 export const createAsyncCallerFactory = asyncTrpc.createCallerFactory;

src/libs/trpc/lambda/index.ts

@@ -0,0 +1 @@
+export * from './serverDatabase';

src/libs/trpc/lambda/serverDatabase.ts

@@ -0,0 +1,10 @@
+import { getServerDB } from '@/database/core/db-adaptor';
+import { trpc } from '@/libs/trpc/init';
+
+export const serverDatabase = trpc.middleware(async (opts) => {
+  const serverDB = await getServerDB();
+
+  return opts.next({
+    ctx: { serverDB },
+  });
+});

src/server/routers/async/file.ts

@@ -11,7 +11,6 @@ import { ChunkModel } from '@/database/models/chunk';
 import { EmbeddingModel } from '@/database/models/embedding';
 import { FileModel } from '@/database/models/file';
 import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
-import { serverDB } from '@/database/server';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
 import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
@@ -31,11 +30,11 @@ const fileProcedure = asyncAuthedProcedure.use(async (opts) => {
 
   return opts.next({
     ctx: {
-      asyncTaskModel: new AsyncTaskModel(serverDB, ctx.userId),
-      chunkModel: new ChunkModel(serverDB, ctx.userId),
+      asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
+      chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
       chunkService: new ChunkService(ctx.userId),
-      embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
-      fileModel: new FileModel(serverDB, ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
     },
   });
 });

src/server/routers/async/ragEval.ts

@@ -7,7 +7,6 @@ import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
 import { ChunkModel } from '@/database/models/chunk';
 import { EmbeddingModel } from '@/database/models/embedding';
 import { FileModel } from '@/database/models/file';
-import { serverDB } from '@/database/server';
 import {
   EvalDatasetRecordModel,
   EvalEvaluationModel,
@@ -25,13 +24,13 @@ const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
 
   return opts.next({
     ctx: {
-      chunkModel: new ChunkModel(serverDB, ctx.userId),
+      chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
       chunkService: new ChunkService(ctx.userId),
       datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
-      embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
       evalRecordModel: new EvaluationRecordModel(ctx.userId),
       evaluationModel: new EvalEvaluationModel(ctx.userId),
-      fileModel: new FileModel(serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
     },
   });
 });