lobehub · arvinxx · Jan 24, 2025
diff --git a/src/app/(main)/settings/storage/page.tsx b/src/app/(main)/settings/storage/page.tsx
@@ -0,0 +1,140 @@
+'use client';
+
+import { ActionIcon } from '@lobehub/ui';
+import { useRequest } from 'ahooks';
+import { Card, Progress, Statistic, Typography, theme } from 'antd';
+import { createStyles } from 'antd-style';
+import { AlertCircle, HardDrive, RotateCw } from 'lucide-react';
+import { Center, Flexbox } from 'react-layout-kit';
+
+import { formatSize } from '@/utils/format';
+
+const { Text } = Typography;
+
+const useStyles = createStyles(({ css, token }) => ({
+  card: css`
+    width: 100%;
+    max-width: 1024px;
+    border: 1px solid ${token.colorBorder};
+    border-radius: ${token.borderRadiusLG}px;
+
+    background: ${token.colorBgContainer};
+    box-shadow: 0 2px 8px ${token.boxShadow};
+
+    .ant-card-body {
+      padding: 24px;
+    }
+  `,
+  detailItem: css`
+    .ant-typography {
+      margin-block-end: 4px;
+
+      &:last-child {
+        margin-block-end: 0;
+      }
+    }
+  `,
+  icon: css`
+    color: ${token.colorPrimary};
+  `,
+  percentage: css`
+    font-size: 24px;
+    font-weight: 600;
+    line-height: 1;
+    color: ${token.colorTextBase};
+  `,
+  progressInfo: css`
+    position: absolute;
+    inset-block-start: 50%;
+    inset-inline-start: 50%;
+    transform: translate(-50%, -50%);
+
+    text-align: center;
+  `,
+  progressWrapper: css`
+    position: relative;
+    width: 180px;
+    height: 180px;
+  `,
+  title: css`
+    margin-block-end: 0;
+    font-size: 16px;
+    font-weight: 500;
+    color: ${token.colorTextBase};
+  `,
+  usageText: css`
+    font-size: 13px;
+    color: ${token.colorTextSecondary};
+  `,
+  warning: css`
+    font-size: 13px;
+    color: ${token.colorWarning};
+  `,
+}));
+
+// 字节转换函数
+
+const StorageEstimate = () => {
+  const { styles } = useStyles();
+  const { token } = theme.useToken();
+
+  const { data, loading, refresh } = useRequest(async () => {
+    const estimate = await navigator.storage.estimate();
+    return {
+      quota: estimate.quota || 0,
+      usage: estimate.usage || 0,
+    };
+  });
+
+  if (!data) return null;
+
+  const usedPercentage = Math.round((data.usage / data.quota) * 100);
+  const freeSpace = data.quota - data.usage;
+  const isLowStorage = usedPercentage > 90;
+
+  return (
+    <Center>
+      <Card
+        className={styles.card}
+        extra={<ActionIcon icon={RotateCw} loading={loading} onClick={refresh} title="Refresh" />}
+        title={
+          <Flexbox align="center" gap={8} horizontal>
+            <HardDrive className={styles.icon} size={18} />
+            <span className={styles.title}>Storage Usage</span>
+          </Flexbox>
+        }
+      >
+        <Flexbox align="center" gap={80} horizontal justify={'center'}>
+          {/* 左侧环形进度区 */}
+          <div className={styles.progressWrapper}>
+            <Progress
+              percent={usedPercentage < 1 ? 1 : usedPercentage}
+              size={180}
+              strokeColor={isLowStorage ? token.colorWarning : token.colorPrimary}
+              strokeWidth={8}
+              type="circle"
+            />
+          </div>
+
+          {/* 右侧详细信息区 */}
+          <Flexbox gap={24}>
+            <Statistic title={'Used Storage'} value={formatSize(data.usage)} />
+            <Statistic title={'Available Storage'} value={formatSize(freeSpace)} />
+            <Statistic title={'Total Storage'} value={formatSize(data.quota)} />
+            {/* 警告信息 */}
+            {isLowStorage && (
+              <Flexbox align="center" gap={8} horizontal>
+                <AlertCircle className={styles.warning} size={16} />
+                <Text className={styles.warning}>
+                  Storage space is running low ({'<'}10% available)
+                </Text>
+              </Flexbox>
+            )}
+          </Flexbox>
+        </Flexbox>
+      </Card>
+    </Center>
+  );
+};
+
+export default StorageEstimate;
diff --git a/src/database/repositories/dataExporter/index.ts b/src/database/repositories/dataExporter/index.ts
@@ -0,0 +1,219 @@
+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;
+}
+
+interface RelationTableConfig {
+  relations: {
+    field: string;
+    sourceField?: string;
+    sourceTable: string;
+  }[];
+  table: keyof typeof EXPORT_TABLES;
+  type: 'relation';
+}
+
+// 配置拆分为基础表和关联表
+
+export const DATA_EXPORT_CONFIG = {
+  // 1. 基础表
+  baseTables: [
+    { table: 'users', userField: 'id' },
+    { table: 'userSettings', userField: 'id' },
+    { table: 'userInstalledPlugins' },
+    { table: 'agents' },
+    { table: 'sessionGroups' },
+    { table: 'sessions' },
+    { table: 'topics' },
+    { table: 'threads' },
+    { table: 'messages' },
+    { table: 'files' },
+    { table: 'knowledgeBases' },
+    { table: 'agentsKnowledgeBases' },
+    { table: 'aiProviders' },
+    { table: 'aiModels' },
+    { table: 'asyncTasks' },
+    { table: 'chunks' },
+    { table: 'embeddings' },
+  ] as BaseTableConfig[],
+  // 2. 关联表
+  relationTables: [
+    {
+      relations: [
+        { sourceField: 'id', sourceTable: 'agents', field: 'agentId' },
+        { sourceField: 'sessionId', sourceTable: 'sessions' },
+      ],
+      table: 'agentsToSessions',
+    },
+    // {
+    //   relations: [
+    //     { sourceField: 'agentId', sourceTable: 'agents' },
+    //     { sourceField: 'fileId', sourceTable: 'files' },
+    //   ],
+    //   table: 'agentsFiles',
+    // },
+    // {
+    //   relations: [{ field: 'sessionId', sourceTable: 'sessions' }],
+    //   table: 'filesToSessions',
+    // },
+    // {
+    //   relations: [{ field: 'id', sourceTable: 'chunks' }],
+    //   table: 'fileChunks',
+    // },
+    {
+      relations: [{ field: 'id', sourceTable: 'messages' }],
+      table: 'messagePlugins',
+    },
+    {
+      relations: [{ field: 'id', sourceTable: 'messages' }],
+      table: 'messageTTS',
+    },
+    {
+      relations: [{ field: 'id', sourceTable: 'messages' }],
+      table: 'messageTranslates',
+    },
+    {
+      relations: [
+        { field: 'messageId', sourceTable: 'messages' },
+        { field: 'fileId', sourceTable: 'files' },
+      ],
+      table: 'messagesFiles',
+    },
+    {
+      relations: [{ field: 'messageId', sourceTable: 'messages' }],
+      table: 'messageQueries',
+    },
+    {
+      relations: [
+        { field: 'messageId', sourceTable: 'messages' },
+        { field: 'chunkId', sourceTable: 'chunks' },
+      ],
+      table: 'messageQueryChunks',
+    },
+    {
+      relations: [
+        { field: 'messageId', sourceTable: 'messages' },
+        { field: 'chunkId', sourceTable: 'chunks' },
+      ],
+      table: 'messageChunks',
+    },
+    {
+      relations: [
+        { field: 'knowledgeBaseId', sourceTable: 'knowledgeBases' },
+        { field: 'fileId', sourceTable: 'files' },
+      ],
+      table: 'knowledgeBaseFiles',
+    },
+  ] as RelationTableConfig[],
+};
+
+export class DataExporterRepos {
+  constructor(
+    private db: LobeChatDatabase,
+    private userId: string,
+  ) {}
+
+  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 {
+      let where;
+
+      const conditions = config.relations.map((relation) => {
+        const sourceData = existingData[relation.sourceTable] || [];
+
+        const sourceIds = sourceData.map((item) => item[relation.sourceField || 'id']);
+        console.log(sourceIds);
+        return inArray(tableObj[relation.field], sourceIds);
+      });
+
+      where = conditions.length === 1 ? conditions[0] : and(...conditions);
+
+      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);
+
+      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;
+    });
+
+    console.log('baseResults:', baseResults);
+    // 2. 然后并发查询所有关联表
+
+    console.log('Querying relation tables...');
+    const relationResults = await pMap(
+      DATA_EXPORT_CONFIG.relationTables,
+      async (config) => ({
+        data: await this.queryTable(config, result),
+        table: config.table,
+      }),
+      { concurrency },
+    );
+
+    // 更新结果集
+    relationResults.forEach(({ table, data }) => {
+      result[table] = data;
+    });
+
+    return result;
+  }
+}
diff --git a/src/database/schemas/user.ts b/src/database/schemas/user.ts
@@ -51,7 +51,7 @@ export const userSettings = pgTable('user_settings', {
 });
 export type UserSettingsItem = typeof userSettings.$inferSelect;
 
-export const installedPlugins = pgTable(
+export const userInstalledPlugins = pgTable(
   'user_installed_plugins',
   {
     userId: text('user_id')
@@ -71,5 +71,5 @@ export const installedPlugins = pgTable(
   }),
 );
 
-export type NewInstalledPlugin = typeof installedPlugins.$inferInsert;
-export type InstalledPluginItem = typeof installedPlugins.$inferSelect;
+export type NewInstalledPlugin = typeof userInstalledPlugins.$inferInsert;
+export type InstalledPluginItem = typeof userInstalledPlugins.$inferSelect;