Skip to content

Commit ffd0dbc

Browse files
authored
♻️ 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
1 parent 3a52f5c commit ffd0dbc

File tree

29 files changed

+174
-112
lines changed

29 files changed

+174
-112
lines changed

src/database/core/db-adaptor.ts

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,24 @@
1-
// import { isDesktop } from '@/const/version';
21
import { getDBInstance } from '@/database/core/web-server';
2+
import { LobeChatDatabase } from '@/database/type';
33

4-
// import { getPgliteInstance } from './electron';
4+
/**
5+
* 懒加载数据库实例
6+
* 避免每次模块导入时都初始化数据库
7+
*/
8+
let cachedDB: LobeChatDatabase | null = null;
9+
10+
export const getServerDB = async (): Promise<LobeChatDatabase> => {
11+
// 如果已经有缓存的实例,直接返回
12+
if (cachedDB) return cachedDB;
13+
14+
try {
15+
// 根据环境选择合适的数据库实例
16+
cachedDB = getDBInstance();
17+
return cachedDB;
18+
} catch (error) {
19+
console.error('❌ Failed to initialize database:', error);
20+
throw error;
21+
}
22+
};
523

624
export const serverDB = getDBInstance();

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
44

55
import { LobeChatDatabase } from '@/database/type';
66
import { ModelProvider } from '@/libs/agent-runtime';
7+
import { sleep } from '@/utils/sleep';
78

89
import { aiProviders, users } from '../../schemas';
910
import { AiProviderModel } from '../aiProvider';
@@ -96,6 +97,7 @@ describe('AiProviderModel', () => {
9697
describe('query', () => {
9798
it('should query ai providers for the user', async () => {
9899
await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
100+
await sleep(10);
99101
await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
100102

101103
const userGroups = await aiProviderModel.query();

src/database/models/session.ts

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,14 @@ export class SessionModel {
193193
type: 'agent' | 'group';
194194
}): Promise<SessionItem> => {
195195
return this.db.transaction(async (trx) => {
196+
if (slug) {
197+
const existResult = await trx.query.sessions.findFirst({
198+
where: and(eq(sessions.slug, slug), eq(sessions.userId, this.userId)),
199+
});
200+
201+
if (existResult) return existResult;
202+
}
203+
196204
const newAgents = await trx
197205
.insert(agents)
198206
.values({

src/database/models/user.ts

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,13 @@ export class UserModel {
8585
const state = result[0];
8686

8787
// Decrypt keyVaults
88-
const decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
88+
let decryptKeyVaults = {};
89+
90+
try {
91+
decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
92+
} catch {
93+
/* empty */
94+
}
8995

9096
const settings: DeepPartial<UserSettings> = {
9197
defaultAgent: state.settingsDefaultAgent || {},

src/database/server/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
export { serverDB } from '../core/db-adaptor';
1+
export { getServerDB, serverDB } from '../core/db-adaptor';

src/libs/trpc/async/index.ts

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,20 @@
1+
import { getServerDB } from '@/database/core/db-adaptor';
2+
13
import { asyncAuth } from './asyncAuth';
24
import { asyncTrpc } from './init';
35

46
export const publicProcedure = asyncTrpc.procedure;
57

68
export const asyncRouter = asyncTrpc.router;
79

8-
export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth);
10+
export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth).use(
11+
asyncTrpc.middleware(async (opts) => {
12+
const serverDB = await getServerDB();
13+
14+
return opts.next({
15+
ctx: { serverDB },
16+
});
17+
}),
18+
);
919

1020
export const createAsyncCallerFactory = asyncTrpc.createCallerFactory;

src/libs/trpc/lambda/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
export * from './serverDatabase';
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
import { getServerDB } from '@/database/core/db-adaptor';
2+
import { trpc } from '@/libs/trpc/init';
3+
4+
export const serverDatabase = trpc.middleware(async (opts) => {
5+
const serverDB = await getServerDB();
6+
7+
return opts.next({
8+
ctx: { serverDB },
9+
});
10+
});

src/server/routers/async/file.ts

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@ import { ChunkModel } from '@/database/models/chunk';
1111
import { EmbeddingModel } from '@/database/models/embedding';
1212
import { FileModel } from '@/database/models/file';
1313
import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
14-
import { serverDB } from '@/database/server';
1514
import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
1615
import { getServerDefaultFilesConfig } from '@/server/globalConfig';
1716
import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
@@ -31,11 +30,11 @@ const fileProcedure = asyncAuthedProcedure.use(async (opts) => {
3130

3231
return opts.next({
3332
ctx: {
34-
asyncTaskModel: new AsyncTaskModel(serverDB, ctx.userId),
35-
chunkModel: new ChunkModel(serverDB, ctx.userId),
33+
asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
34+
chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
3635
chunkService: new ChunkService(ctx.userId),
37-
embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
38-
fileModel: new FileModel(serverDB, ctx.userId),
36+
embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
37+
fileModel: new FileModel(ctx.serverDB, ctx.userId),
3938
},
4039
});
4140
});

src/server/routers/async/ragEval.ts

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
77
import { ChunkModel } from '@/database/models/chunk';
88
import { EmbeddingModel } from '@/database/models/embedding';
99
import { FileModel } from '@/database/models/file';
10-
import { serverDB } from '@/database/server';
1110
import {
1211
EvalDatasetRecordModel,
1312
EvalEvaluationModel,
@@ -25,13 +24,13 @@ const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
2524

2625
return opts.next({
2726
ctx: {
28-
chunkModel: new ChunkModel(serverDB, ctx.userId),
27+
chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
2928
chunkService: new ChunkService(ctx.userId),
3029
datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
31-
embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
30+
embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
3231
evalRecordModel: new EvaluationRecordModel(ctx.userId),
3332
evaluationModel: new EvalEvaluationModel(ctx.userId),
34-
fileModel: new FileModel(serverDB, ctx.userId),
33+
fileModel: new FileModel(ctx.serverDB, ctx.userId),
3534
},
3635
});
3736
});

0 commit comments

Comments
 (0)