♻️ refactor: refactor database file model to remove server env (lobeh…

…ub#4990)

* refactor the db model

* Update file.ts

* fix test

src/config/db.ts

@@ -11,20 +11,20 @@ export const getServerDBConfig = () => {
       DATABASE_TEST_URL: process.env.DATABASE_TEST_URL,
       DATABASE_URL: process.env.DATABASE_URL,
 
-      DISABLE_REMOVE_GLOBAL_FILE: process.env.DISABLE_REMOVE_GLOBAL_FILE === '1',
-
       KEY_VAULTS_SECRET: process.env.KEY_VAULTS_SECRET,
 
       NEXT_PUBLIC_ENABLED_SERVER_SERVICE: process.env.NEXT_PUBLIC_SERVICE_MODE === 'server',
+
+      REMOVE_GLOBAL_FILE: process.env.DISABLE_REMOVE_GLOBAL_FILE !== '0',
     },
     server: {
       DATABASE_DRIVER: z.enum(['neon', 'node']),
       DATABASE_TEST_URL: z.string().optional(),
       DATABASE_URL: z.string().optional(),
 
-      DISABLE_REMOVE_GLOBAL_FILE: z.boolean().optional(),
-
       KEY_VAULTS_SECRET: z.string().optional(),
+
+      REMOVE_GLOBAL_FILE: z.boolean().optional(),
     },
   });
 };

src/database/server/models/__tests__/file.test.ts

@@ -8,21 +8,7 @@ import { FilesTabs, SortType } from '@/types/files';
 import { files, globalFiles, knowledgeBaseFiles, knowledgeBases, users } from '../../../schemas';
 import { FileModel } from '../file';
 
-let serverDB = await getTestDBInstance();
-
-let DISABLE_REMOVE_GLOBAL_FILE = false;
-
-vi.mock('@/config/db', async () => ({
-  get serverDBEnv() {
-    return {
-      get DISABLE_REMOVE_GLOBAL_FILE() {
-        return DISABLE_REMOVE_GLOBAL_FILE;
-      },
-      DATABASE_TEST_URL: process.env.DATABASE_TEST_URL,
-      DATABASE_DRIVER: 'node',
-    };
-  },
-}));
+const serverDB = await getTestDBInstance();
 
 const userId = 'file-model-test-user-id';
 const fileModel = new FileModel(serverDB, userId);
@@ -146,7 +132,6 @@ describe('FileModel', () => {
       expect(globalFile).toBeUndefined();
     });
     it('should delete a file by id but global file not removed ', async () => {
-      DISABLE_REMOVE_GLOBAL_FILE = true;
       await fileModel.createGlobalFile({
         hashId: '1',
         url: 'https://example.com/file1.txt',
@@ -162,7 +147,7 @@ describe('FileModel', () => {
         fileHash: '1',
       });
 
-      await fileModel.delete(id);
+      await fileModel.delete(id, false);
 
       const file = await serverDB.query.files.findFirst({ where: eq(files.id, id) });
       const globalFile = await serverDB.query.globalFiles.findFirst({
@@ -171,7 +156,6 @@ describe('FileModel', () => {
 
       expect(file).toBeUndefined();
       expect(globalFile).toBeDefined();
-      DISABLE_REMOVE_GLOBAL_FILE = false;
     });
   });
 
@@ -225,7 +209,6 @@ describe('FileModel', () => {
       expect(globalFilesResult2).toHaveLength(0);
     });
     it('should delete multiple files but not remove global files if DISABLE_REMOVE_GLOBAL_FILE=true', async () => {
-      DISABLE_REMOVE_GLOBAL_FILE = true;
       await fileModel.createGlobalFile({
         hashId: '1',
         url: 'https://example.com/file1.txt',
@@ -260,7 +243,7 @@ describe('FileModel', () => {
 
       expect(globalFilesResult).toHaveLength(2);
 
-      await fileModel.deleteMany([file1.id, file2.id]);
+      await fileModel.deleteMany([file1.id, file2.id], false);
 
       const remainingFiles = await serverDB.query.files.findMany({
         where: eq(files.userId, userId),
@@ -271,7 +254,6 @@ describe('FileModel', () => {
 
       expect(remainingFiles).toHaveLength(0);
       expect(globalFilesResult2).toHaveLength(2);
-      DISABLE_REMOVE_GLOBAL_FILE = false;
     });
   });
 

src/database/server/models/_template.ts

@@ -1,5 +1,4 @@
-import { eq } from 'drizzle-orm';
-import { and, desc } from 'drizzle-orm/expressions';
+import { and, desc, eq } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 

src/database/server/models/agent.ts

@@ -1,7 +1,4 @@
-import { inArray } from 'drizzle-orm';
-import { and, desc, eq } from 'drizzle-orm/expressions';
-
-import { LobeChatDatabase } from '@/database/type';
+import { and, desc, eq, inArray } from 'drizzle-orm/expressions';
 
 import {
   agents,
@@ -11,6 +8,7 @@ import {
   files,
   knowledgeBases,
 } from '@/database/schemas';
+import { LobeChatDatabase } from '@/database/type';
 
 export class AgentModel {
   private userId: string;

src/database/server/models/asyncTask.ts

@@ -1,5 +1,4 @@
-import { eq, inArray, lt } from 'drizzle-orm';
-import { and } from 'drizzle-orm/expressions';
+import { and, eq, inArray, lt } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 import {

src/database/server/models/chunk.ts

@@ -1,5 +1,5 @@
-import { asc, cosineDistance, count, eq, inArray, sql } from 'drizzle-orm';
-import { and, desc, isNull } from 'drizzle-orm/expressions';
+import { cosineDistance, count, sql } from 'drizzle-orm';
+import { and, asc, desc, eq, inArray, isNull } from 'drizzle-orm/expressions';
 import { chunk } from 'lodash-es';
 
 import { LobeChatDatabase } from '@/database/type';

src/database/server/models/embedding.ts

@@ -1,5 +1,5 @@
-import { count, eq } from 'drizzle-orm';
-import { and } from 'drizzle-orm/expressions';
+import { count } from 'drizzle-orm';
+import { and, eq } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 

src/database/server/models/file.ts

@@ -2,7 +2,6 @@ import { asc, count, eq, ilike, inArray, notExists, or, sum } from 'drizzle-orm'
 import { and, desc, like } from 'drizzle-orm/expressions';
 import type { PgTransaction } from 'drizzle-orm/pg-core';
 
-import { serverDBEnv } from '@/config/db';
 import { LobeChatDatabase } from '@/database/type';
 import { FilesTabs, QueryFileListParams, SortType } from '@/types/files';
 
@@ -67,7 +66,7 @@ export class FileModel {
     };
   };
 
-  delete = async (id: string) => {
+  delete = async (id: string, removeGlobalFile: boolean = true) => {
     const file = await this.findById(id);
     if (!file) return;
 
@@ -89,7 +88,7 @@ export class FileModel {
 
       // delete the file from global file if it is not used by other files
       // if `DISABLE_REMOVE_GLOBAL_FILE` is true, we will not remove the global file
-      if (fileCount === 0 && !serverDBEnv.DISABLE_REMOVE_GLOBAL_FILE) {
+      if (fileCount === 0 && removeGlobalFile) {
         await trx.delete(globalFiles).where(eq(globalFiles.hashId, fileHash));
 
         return file;
@@ -112,7 +111,7 @@ export class FileModel {
     return parseInt(result[0].totalSize!) || 0;
   };
 
-  deleteMany = async (ids: string[]) => {
+  deleteMany = async (ids: string[], removeGlobalFile: boolean = true) => {
     const fileList = await this.findByIds(ids);
     const hashList = fileList.map((file) => file.fileHash!);
 
@@ -144,7 +143,7 @@ export class FileModel {
 
       const needToDeleteList = fileHashCounts.filter((item) => item.count === 0);
 
-      if (needToDeleteList.length === 0 || serverDBEnv.DISABLE_REMOVE_GLOBAL_FILE) return;
+      if (needToDeleteList.length === 0 || !removeGlobalFile) return;
 
       // delete the file from global file if it is not used by other files
       await trx.delete(globalFiles).where(

src/database/server/models/knowledgeBase.ts

@@ -1,5 +1,4 @@
-import { eq, inArray } from 'drizzle-orm';
-import { and, desc } from 'drizzle-orm/expressions';
+import { and, desc, eq, inArray } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 import { KnowledgeBaseItem } from '@/types/knowledgeBase';

src/database/server/models/message.ts

@@ -1,8 +1,8 @@
 import { count } from 'drizzle-orm';
 import { and, asc, desc, eq, gte, inArray, isNull, like, lt } from 'drizzle-orm/expressions';
 
-import { idGenerator } from '@/database/utils/idGenerator';
 import { LobeChatDatabase } from '@/database/type';
+import { idGenerator } from '@/database/utils/idGenerator';
 import { getFullFileUrl } from '@/server/utils/files';
 import {
   ChatFileItem,

src/database/server/models/session.ts

@@ -1,5 +1,5 @@
-import { Column, asc, count, inArray, like, sql } from 'drizzle-orm';
-import { and, desc, eq, not, or } from 'drizzle-orm/expressions';
+import { Column, count, sql } from 'drizzle-orm';
+import { and, asc, desc, eq, inArray, like, not, or } from 'drizzle-orm/expressions';
 
 import { appEnv } from '@/config/app';
 import { INBOX_SESSION_ID } from '@/const/session';

src/database/server/models/sessionGroup.ts

@@ -1,5 +1,4 @@
-import { eq } from 'drizzle-orm';
-import { and, asc, desc } from 'drizzle-orm/expressions';
+import { and, asc, desc, eq } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 import { idGenerator } from '@/database/utils/idGenerator';

src/database/server/models/thread.ts

@@ -1,5 +1,4 @@
-import { eq } from 'drizzle-orm';
-import { and, desc } from 'drizzle-orm/expressions';
+import { and, desc, eq } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 import { CreateThreadParams, ThreadStatus } from '@/types/topic';

src/database/server/models/topic.ts

@@ -1,5 +1,5 @@
-import { Column, count, inArray, sql } from 'drizzle-orm';
-import { and, desc, eq, exists, isNull, like, or } from 'drizzle-orm/expressions';
+import { Column, count, sql } from 'drizzle-orm';
+import { and, desc, eq, exists, inArray, isNull, like, or } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
 import { idGenerator } from '@/database/utils/idGenerator';

src/database/server/models/user.ts

@@ -1,5 +1,5 @@
 import { TRPCError } from '@trpc/server';
-import { eq } from 'drizzle-orm';
+import { eq } from 'drizzle-orm/expressions';
 import { DeepPartial } from 'utility-types';
 
 import { LobeChatDatabase } from '@/database/type';

src/server/routers/async/file.ts

@@ -3,14 +3,15 @@ import { chunk } from 'lodash-es';
 import pMap from 'p-map';
 import { z } from 'zod';
 
+import { serverDBEnv } from '@/config/db';
 import { fileEnv } from '@/config/file';
 import { DEFAULT_EMBEDDING_MODEL } from '@/const/settings';
+import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
 import { serverDB } from '@/database/server';
 import { ASYNC_TASK_TIMEOUT, AsyncTaskModel } from '@/database/server/models/asyncTask';
 import { ChunkModel } from '@/database/server/models/chunk';
 import { EmbeddingModel } from '@/database/server/models/embedding';
 import { FileModel } from '@/database/server/models/file';
-import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
 import { ModelProvider } from '@/libs/agent-runtime';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
 import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
@@ -175,7 +176,7 @@ export const fileRouter = router({
         console.error(e);
         // if file not found, delete it from db
         if ((e as any).Code === 'NoSuchKey') {
-          await ctx.fileModel.delete(input.fileId);
+          await ctx.fileModel.delete(input.fileId, serverDBEnv.REMOVE_GLOBAL_FILE);
           throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
         }
       }

src/server/routers/lambda/file.ts

@@ -1,6 +1,7 @@
 import { TRPCError } from '@trpc/server';
 import { z } from 'zod';
 
+import { serverDBEnv } from '@/config/db';
 import { serverDB } from '@/database/server';
 import { AsyncTaskModel } from '@/database/server/models/asyncTask';
 import { ChunkModel } from '@/database/server/models/chunk';
@@ -153,7 +154,7 @@ export const fileRouter = router({
   }),
 
   removeFile: fileProcedure.input(z.object({ id: z.string() })).mutation(async ({ input, ctx }) => {
-    const file = await ctx.fileModel.delete(input.id);
+    const file = await ctx.fileModel.delete(input.id, serverDBEnv.REMOVE_GLOBAL_FILE);
 
     if (!file) return;
 
@@ -184,7 +185,10 @@ export const fileRouter = router({
   removeFiles: fileProcedure
     .input(z.object({ ids: z.array(z.string()) }))
     .mutation(async ({ input, ctx }) => {
-      const needToRemoveFileList = await ctx.fileModel.deleteMany(input.ids);
+      const needToRemoveFileList = await ctx.fileModel.deleteMany(
+        input.ids,
+        serverDBEnv.REMOVE_GLOBAL_FILE,
+      );
 
       if (!needToRemoveFileList || needToRemoveFileList.length === 0) return;