@lobehub/chat 1.76.0 → 1.77.0

package/src/database/repositories/dataExporter/index.test.ts

@@ -0,0 +1,330 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { clientDB, initializeDB } from '@/database/client/db';
+import {
+  agents,
+  agentsKnowledgeBases,
+  agentsToSessions,
+  files,
+  filesToSessions,
+  globalFiles,
+  knowledgeBaseFiles,
+  knowledgeBases,
+  messages,
+  sessionGroups,
+  sessions,
+  topics,
+  userSettings,
+  users,
+} from '@/database/schemas';
+import { LobeChatDatabase } from '@/database/type';
+
+import { DATA_EXPORT_CONFIG, DataExporterRepos } from './index';
+
+let db = clientDB as LobeChatDatabase;
+
+// 设置测试数据
+describe('DataExporterRepos', () => {
+  // 测试数据 ID
+  const testIds = {
+    userId: 'test-user-id',
+    fileId: 'test-file-id',
+    fileHash: 'test-file-hash',
+    sessionId: 'test-session-id',
+    agentId: 'test-agent-id',
+    topicId: 'test-topic-id',
+    messageId: 'test-message-id',
+    knowledgeBaseId: 'test-kb-id',
+  };
+
+  // 设置测试环境
+  let userId: string = testIds.userId;
+
+  const setupTestData = async () => {
+    await db.transaction(async (trx) => {
+      // 用户数据
+      await trx.insert(users).values({
+        id: testIds.userId,
+        username: 'testuser',
+        email: 'test@example.com',
+      });
+
+      // 用户设置
+      await trx.insert(userSettings).values({
+        id: testIds.userId,
+        general: { theme: 'light' },
+      });
+
+      // 全局文件
+      await trx.insert(globalFiles).values({
+        hashId: testIds.fileHash,
+        fileType: 'text/plain',
+        size: 1024,
+        url: 'https://example.com/test-file.txt',
+        creator: testIds.userId,
+      });
+
+      // 文件数据
+      await trx.insert(files).values({
+        id: testIds.fileId,
+        userId: testIds.userId,
+        fileType: 'text/plain',
+        fileHash: testIds.fileHash,
+        name: 'test-file.txt',
+        size: 1024,
+        url: 'https://example.com/test-file.txt',
+      });
+
+      // 会话组
+      await trx.insert(sessionGroups).values({
+        name: 'Test Group',
+        userId: testIds.userId,
+      });
+
+      // 会话
+      await trx.insert(sessions).values({
+        id: testIds.sessionId,
+        slug: 'test-session',
+        title: 'Test Session',
+        userId: testIds.userId,
+      });
+
+      // 主题
+      await trx.insert(topics).values({
+        id: testIds.topicId,
+        title: 'Test Topic',
+        sessionId: testIds.sessionId,
+        userId: testIds.userId,
+      });
+
+      // 消息
+      await trx.insert(messages).values({
+        id: testIds.messageId,
+        role: 'user',
+        content: 'Hello, world!',
+        userId: testIds.userId,
+        sessionId: testIds.sessionId,
+        topicId: testIds.topicId,
+      });
+
+      // 代理
+      await trx.insert(agents).values({
+        id: testIds.agentId,
+        title: 'Test Agent',
+        userId: testIds.userId,
+      });
+
+      // 代理到会话的关联
+      await trx.insert(agentsToSessions).values({
+        agentId: testIds.agentId,
+        sessionId: testIds.sessionId,
+        userId: testIds.userId,
+      });
+
+      // 文件到会话的关联
+      await trx.insert(filesToSessions).values({
+        fileId: testIds.fileId,
+        sessionId: testIds.sessionId,
+        userId: testIds.userId,
+      });
+
+      // 知识库
+      await trx.insert(knowledgeBases).values({
+        id: testIds.knowledgeBaseId,
+        name: 'Test Knowledge Base',
+        userId: testIds.userId,
+      });
+
+      // 知识库文件
+      await trx.insert(knowledgeBaseFiles).values({
+        knowledgeBaseId: testIds.knowledgeBaseId,
+        fileId: testIds.fileId,
+        userId: testIds.userId,
+      });
+
+      // 代理知识库
+      await trx.insert(agentsKnowledgeBases).values({
+        agentId: testIds.agentId,
+        knowledgeBaseId: testIds.knowledgeBaseId,
+        userId: testIds.userId,
+      });
+    });
+  };
+
+  beforeEach(async () => {
+    // 创建内存数据库
+    await initializeDB();
+
+    // 插入测试数据
+    await setupTestData();
+  });
+
+  afterEach(async () => {
+    await db.delete(users);
+    await db.delete(globalFiles);
+
+    vi.restoreAllMocks();
+  });
+
+  describe('export', () => {
+    it('should export all user data correctly', async () => {
+      // 创建导出器实例
+      const dataExporter = new DataExporterRepos(db, userId);
+
+      // 执行导出
+      const result = await dataExporter.export();
+
+      // 验证基础表导出结果
+      // expect(result).toHaveProperty('users');
+      // expect(result.users).toHaveLength(1);
+      // expect(result.users[0]).toHaveProperty('id', testIds.userId);
+      // expect(result.users[0]).not.toHaveProperty('userId'); // userId 字段应该被移除
+
+      expect(result).toHaveProperty('userSettings');
+      expect(result.userSettings).toHaveLength(1);
+      expect(result.userSettings[0]).toHaveProperty('id', testIds.userId);
+
+      // expect(result).toHaveProperty('files');
+      // expect(result.files).toHaveLength(1);
+      // expect(result.files[0]).toHaveProperty('id', testIds.fileId);
+      // expect(result.files[0]).toHaveProperty('fileHash', testIds.fileHash);
+      // expect(result.files[0]).not.toHaveProperty('userId');
+
+      expect(result).toHaveProperty('sessions');
+      expect(result.sessions).toHaveLength(1);
+      expect(result.sessions[0]).toHaveProperty('id', testIds.sessionId);
+
+      expect(result).toHaveProperty('topics');
+      expect(result.topics).toHaveLength(1);
+      expect(result.topics[0]).toHaveProperty('id', testIds.topicId);
+
+      expect(result).toHaveProperty('messages');
+      expect(result.messages).toHaveLength(1);
+      expect(result.messages[0]).toHaveProperty('id', testIds.messageId);
+
+      expect(result).toHaveProperty('agents');
+      expect(result.agents).toHaveLength(1);
+      expect(result.agents[0]).toHaveProperty('id', testIds.agentId);
+
+      // expect(result).toHaveProperty('knowledgeBases');
+      // expect(result.knowledgeBases).toHaveLength(1);
+      // expect(result.knowledgeBases[0]).toHaveProperty('id', testIds.knowledgeBaseId);
+
+      // 验证关联表导出结果
+      // expect(result).toHaveProperty('globalFiles');
+      // expect(result.globalFiles).toHaveLength(1);
+      // expect(result.globalFiles[0]).toHaveProperty('hashId', testIds.fileHash);
+
+      expect(result).toHaveProperty('agentsToSessions');
+      expect(result.agentsToSessions).toHaveLength(1);
+      expect(result.agentsToSessions[0]).toHaveProperty('agentId', testIds.agentId);
+      expect(result.agentsToSessions[0]).toHaveProperty('sessionId', testIds.sessionId);
+
+      // expect(result).toHaveProperty('filesToSessions');
+      // expect(result.filesToSessions).toHaveLength(1);
+      // expect(result.filesToSessions[0]).toHaveProperty('fileId', testIds.fileId);
+      // expect(result.filesToSessions[0]).toHaveProperty('sessionId', testIds.sessionId);
+
+      // expect(result).toHaveProperty('knowledgeBaseFiles');
+      // expect(result.knowledgeBaseFiles).toHaveLength(1);
+      // expect(result.knowledgeBaseFiles[0]).toHaveProperty(
+      //   'knowledgeBaseId',
+      //   testIds.knowledgeBaseId,
+      // );
+      // expect(result.knowledgeBaseFiles[0]).toHaveProperty('fileId', testIds.fileId);
+    });
+
+    it('should handle empty database gracefully', async () => {
+      // 清空数据库
+
+      await db.delete(users);
+      await db.delete(globalFiles);
+
+      // 创建导出器实例
+      const dataExporter = new DataExporterRepos(db, userId);
+
+      // 执行导出
+      const result = await dataExporter.export();
+
+      // 验证所有表都返回空数组
+      DATA_EXPORT_CONFIG.baseTables.forEach(({ table }) => {
+        expect(result).toHaveProperty(table);
+        expect(result[table]).toEqual([]);
+      });
+
+      DATA_EXPORT_CONFIG.relationTables.forEach(({ table }) => {
+        expect(result).toHaveProperty(table);
+        expect(result[table]).toEqual([]);
+      });
+    });
+
+    it('should handle database query errors', async () => {
+      // 模拟查询错误
+      // @ts-ignore
+      vi.spyOn(db.query.users, 'findMany').mockRejectedValueOnce(new Error('Database error'));
+
+      // 创建导出器实例
+      const dataExporter = new DataExporterRepos(db, userId);
+
+      // 执行导出
+      const result = await dataExporter.export();
+
+      // 验证其他表仍然被导出
+      expect(result).toHaveProperty('sessions');
+      expect(result.sessions).toHaveLength(1);
+    });
+
+    it.skip('should skip relation tables when source tables have no data', async () => {
+      // 删除文件数据，这将导致 globalFiles 表被跳过
+      await db.delete(files);
+
+      // 创建导出器实例
+      const dataExporter = new DataExporterRepos(db, userId);
+
+      // 执行导出
+      const result = await dataExporter.export();
+
+      // 验证文件表为空
+      // expect(result).toHaveProperty('files');
+      // expect(result.files).toEqual([]);
+
+      // 验证关联表也为空
+      // expect(result).toHaveProperty('globalFiles');
+      // expect(result.globalFiles).toEqual([]);
+    });
+
+    it('should export data for a different user', async () => {
+      // 创建另一个用户
+      const anotherUserId = 'another-user-id';
+      await db.transaction(async (trx) => {
+        await trx.insert(users).values({
+          id: anotherUserId,
+          username: 'anotheruser',
+          email: 'another@example.com',
+        });
+        await trx.insert(sessions).values({
+          id: 'another-session-id',
+          slug: 'another-session',
+          title: 'Another Session',
+          userId: anotherUserId,
+        });
+      });
+
+      // 创建导出器实例，使用另一个用户 ID
+      const dataExporter = new DataExporterRepos(db, anotherUserId);
+
+      // 执行导出
+      const result = await dataExporter.export();
+
+      // 验证只导出了另一个用户的数据
+      // expect(result).toHaveProperty('users');
+      // expect(result.users).toHaveLength(1);
+      // expect(result.users[0]).toHaveProperty('id', anotherUserId);
+
+      expect(result).toHaveProperty('sessions');
+      expect(result.sessions).toHaveLength(1);
+      expect(result.sessions[0]).not.toHaveProperty('userId', anotherUserId);
+      expect(result.sessions[0]).toHaveProperty('id', 'another-session-id');
+    });
+  });
+});

package/src/database/repositories/dataExporter/index.ts

@@ -0,0 +1,216 @@
+import { and, eq, inArray } from 'drizzle-orm/expressions';
+import pMap from 'p-map';
+
+import * as EXPORT_TABLES from '@/database/schemas';
+import { LobeChatDatabase } from '@/database/type';
+
+interface BaseTableConfig {
+  table: keyof typeof EXPORT_TABLES;
+  type: 'base';
+  userField?: string;
+}
+
+export interface RelationTableConfig {
+  relations: {
+    field: string;
+    sourceField?: string;
+    sourceTable: keyof typeof EXPORT_TABLES;
+  }[];
+  table: keyof typeof EXPORT_TABLES;
+  type: 'relation';
+}
+
+export const DATA_EXPORT_CONFIG = {
+  baseTables: [
+    // { table: 'users', userField: 'id' },
+    { table: 'userSettings', userField: 'id' },
+    { table: 'userInstalledPlugins' },
+    { table: 'agents' },
+    // { table: 'agentsFiles' },
+    // { table: 'agentsKnowledgeBases' },
+    // { table: 'agentsToSessions' },
+    { table: 'aiModels' },
+    { table: 'aiProviders' },
+    // async tasks should not be included
+    // { table: 'asyncTasks' },
+    // { table: 'chunks' },
+    // { table: 'unstructuredChunks' },
+    // { table: 'embeddings' },
+    // { table: 'files' },
+    // { table: 'fileChunks' },
+    // { table: 'filesToSessions' },
+    // { table: 'knowledgeBases' },
+    // { table: 'knowledgeBaseFiles' },
+    { table: 'messageChunks' },
+    { table: 'messagePlugins' },
+    // { table: 'messageQueryChunks' },
+    // { table: 'messageQueries' },
+    { table: 'messageTranslates' },
+    // { table: 'messageTTS' },
+    { table: 'messages' },
+    // { table: 'messagesFiles' },
+
+    // next auth tables won't be included
+    // { table: 'nextauthAccounts' },
+    // { table: 'nextauthSessions' },
+    // { table: 'nextauthAuthenticators' },
+    // { table: 'nextauthVerificationTokens' },
+    { table: 'sessionGroups' },
+    { table: 'sessions' },
+    { table: 'threads' },
+    { table: 'topics' },
+  ] as BaseTableConfig[],
+  relationTables: [
+    // {
+    //   relations: [{ field: 'hashId', sourceField: 'fileHash', sourceTable: 'files' }],
+    //   table: 'globalFiles',
+    // },
+    {
+      relations: [
+        { field: 'agentId', sourceField: 'id', sourceTable: 'agents' },
+        { field: 'sessionId', sourceField: 'id', sourceTable: 'sessions' },
+      ],
+      table: 'agentsToSessions',
+    },
+
+    // {
+    //   relations: [{ field: 'id', sourceField: 'id', sourceTable: 'messages' }],
+    //   table: 'messagePlugins',
+    // },
+  ] as RelationTableConfig[],
+};
+
+export class DataExporterRepos {
+  private userId: string;
+  private db: LobeChatDatabase;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
+    this.userId = userId;
+  }
+
+  private removeUserId(data: any[]) {
+    return data.map((item) => {
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      const { userId: _, ...rest } = item;
+      return rest;
+    });
+  }
+
+  private async queryTable(config: RelationTableConfig, existingData: Record<string, any[]>) {
+    const { table } = config;
+    const tableObj = EXPORT_TABLES[table];
+    if (!tableObj) throw new Error(`Table ${table} not found`);
+
+    try {
+      const conditions = [];
+
+      // 处理每个关联条件
+      for (const relation of config.relations) {
+        const sourceData = existingData[relation.sourceTable] || [];
+
+        // 如果源数据为空，这个表可能无法查询到任何数据
+        if (sourceData.length === 0) {
+          console.log(
+            `Source table ${relation.sourceTable} has no data, skipping query for ${table}`,
+          );
+          return [];
+        }
+
+        const sourceIds = sourceData.map((item) => item[relation.sourceField || 'id']);
+        conditions.push(inArray(tableObj[relation.field], sourceIds));
+      }
+
+      // 如果表有userId字段并且不是users表，添加用户过滤
+      if ('userId' in tableObj && table !== 'users' && !config.relations) {
+        conditions.push(eq(tableObj.userId, this.userId));
+      }
+
+      // 组合所有条件
+      const where = conditions.length === 1 ? conditions[0] : and(...conditions);
+
+      // @ts-expect-error query
+      const result = await this.db.query[table].findMany({ where });
+
+      // 只对使用 userId 查询的表移除 userId 字段
+      console.log(`Successfully exported table: ${table}, count: ${result.length}`);
+      return config.relations ? result : this.removeUserId(result);
+    } catch (error) {
+      console.error(`Error querying table ${table}:`, error);
+      return [];
+    }
+  }
+
+  private async queryBaseTables(config: BaseTableConfig) {
+    const { table } = config;
+    const tableObj = EXPORT_TABLES[table];
+    if (!tableObj) throw new Error(`Table ${table} not found`);
+
+    try {
+      // 如果有关联配置，使用关联查询
+
+      // 默认使用 userId 查询，特殊情况使用 userField
+      const userField = config.userField || 'userId';
+      const where = eq(tableObj[userField], this.userId);
+
+      // @ts-expect-error query
+      const result = await this.db.query[table].findMany({ where });
+
+      // 只对使用 userId 查询的表移除 userId 字段
+      console.log(`Successfully exported table: ${table}, count: ${result.length}`);
+      return this.removeUserId(result);
+    } catch (error) {
+      console.error(`Error querying table ${table}:`, error);
+      return [];
+    }
+  }
+
+  async export(concurrency = 10) {
+    const result: Record<string, any[]> = {};
+
+    // 1. 首先并发查询所有基础表
+    console.log('Querying base tables...');
+    const baseResults = await pMap(
+      DATA_EXPORT_CONFIG.baseTables,
+      async (config) => ({ data: await this.queryBaseTables(config), table: config.table }),
+      { concurrency },
+    );
+
+    // 更新结果集
+    baseResults.forEach(({ table, data }) => {
+      result[table] = data;
+    });
+
+    // 2. 然后并发查询所有关联表
+
+    const relationResults = await pMap(
+      DATA_EXPORT_CONFIG.relationTables,
+      async (config) => {
+        // 检查所有依赖的源表是否有数据
+        const allSourcesHaveData = config.relations.every(
+          (relation) => (result[relation.sourceTable] || []).length > 0,
+        );
+
+        if (!allSourcesHaveData) {
+          console.log(`Skipping table ${config.table} as some source tables have no data`);
+          return { data: [], table: config.table };
+        }
+
+        return {
+          data: await this.queryTable(config, result),
+          table: config.table,
+        };
+      },
+      { concurrency },
+    );
+
+    // 更新结果集
+    relationResults.forEach(({ table, data }) => {
+      result[table] = data;
+    });
+
+    console.log('finalResults:', result);
+
+    return result;
+  }
+}

package/src/database/repositories/dataImporter/__tests__/fixtures/agents.json

@@ -0,0 +1,65 @@
+{
+  "data": {
+    "agents": [
+      {
+        "id": "agt_CSLkGPIXIQms",
+        "slug": "room-possible-somewhere-numeral",
+        "title": null,
+        "description": null,
+        "tags": [],
+        "avatar": null,
+        "backgroundColor": null,
+        "plugins": [],
+        "chatConfig": {
+          "searchMode": "off"
+        },
+        "fewShots": null,
+        "model": "gemini-2.0-flash-exp",
+        "params": {
+          "top_p": 1,
+          "temperature": 1,
+          "presence_penalty": 0,
+          "frequency_penalty": 0
+        },
+        "provider": "google",
+        "systemRole": "",
+        "tts": {
+          "voice": {
+            "openai": "alloy"
+          },
+          "sttLocale": "auto",
+          "ttsService": "openai",
+          "showAllLocaleVoice": false
+        },
+        "accessedAt": "2025-03-05T05:25:44.521Z",
+        "createdAt": "2025-03-05T05:25:44.522Z",
+        "updatedAt": "2025-03-05T05:25:44.522Z"
+      }
+    ],
+    "agentsToSessions": [
+      {
+        "agentId": "agt_CSLkGPIXIQms",
+        "sessionId": "ssn_WdL6m1mReqyb"
+      }
+    ],
+    "sessions": [
+      {
+        "id": "ssn_WdL6m1mReqyb",
+        "slug": "inbox-t",
+        "title": null,
+        "description": null,
+        "avatar": null,
+        "backgroundColor": null,
+        "type": "agent",
+        "groupId": null,
+        "clientId": null,
+        "pinned": false,
+        "accessedAt": "2025-03-05T05:25:44.521Z",
+        "createdAt": "2025-03-05T05:25:44.528Z",
+        "updatedAt": "2025-03-05T05:25:44.528Z"
+      }
+    ]
+  },
+  "mode": "pglite",
+  "schemaHash": "89e91285be422d5f44511c7405f57b57f8dfda4f0304126ae9b0f266e5fd60f1"
+}