d8d-vite-starter/src/server/utils/restore.ts

import { pgRestore } from 'pg-dump-restore'
import { promises as fs } from 'fs'
import path from 'path'
import { logger } from './logger'
import process from 'node:process'

const BACKUP_DIR = process.env.BACKUP_DIR || './backups'

export interface RestoreOptions {
  host?: string
  port?: number
  database?: string
  username?: string
  password?: string
  backupFile?: string
}

export class DatabaseRestore {
  private getDbConfig() {
    return {
      host: process.env.DB_HOST || 'localhost',
      port: parseInt(process.env.DB_PORT || '5432'),
      database: process.env.DB_DATABASE || 'postgres',
      username: process.env.DB_USERNAME || 'postgres',
      password: process.env.DB_PASSWORD || '',
    }
  }

  async restoreDatabase(options: Partial<RestoreOptions> = {}): Promise<void> {
    const config = { ...this.getDbConfig(), ...options }
    let backupFile = options.backupFile

    if (!backupFile) {
      const latestBackup = await this.findLatestBackup()
      if (!latestBackup) {
        throw new Error('未找到可用的备份文件')
      }
      backupFile = latestBackup
    }

    if (!await this.backupExists(backupFile)) {
      throw new Error(`备份文件不存在: ${backupFile}`)
    }

    try {
      logger.db(`开始恢复数据库备份: ${backupFile}`)

      await pgRestore({
        host: config.host,
        port: config.port,
        database: config.database,
        username: config.username,
        password: config.password,
      }, {
        filePath: backupFile,
      })

      logger.db(`数据库恢复成功: ${config.database}`)
    } catch (error) {
      logger.error(`数据库恢复失败: ${error}`)
      throw error
    }
  }

  async findLatestBackup(): Promise<string | null> {
    try {
      const files = await fs.readdir(BACKUP_DIR)
      const backupFiles = files
        .filter(f => f.endsWith('.dump'))
        .sort()
        .reverse()

      return backupFiles.length > 0 ? path.join(BACKUP_DIR, backupFiles[0]) : null
    } catch (error) {
      logger.error(`查找备份文件失败: ${error}`)
      return null
    }
  }

  async listBackups(): Promise<string[]> {
    try {
      const files = await fs.readdir(BACKUP_DIR)
      return files.filter(f => f.endsWith('.dump'))
    } catch (error) {
      logger.error(`列出备份文件失败: ${error}`)
      return []
    }
  }

  async backupExists(backupFile: string): Promise<boolean> {
    try {
      await fs.access(backupFile)
      return true
    } catch {
      return false
    }
  }

  async getBackupInfo(backupFile: string): Promise<{
    size: number
    mtime: Date
    formattedSize: string
  }> {
    try {
      const stats = await fs.stat(backupFile)
      return {
        size: stats.size,
        mtime: stats.mtime,
        formattedSize: this.formatFileSize(stats.size)
      }
    } catch (error) {
      throw new Error(`获取备份信息失败: ${error}`)
    }
  }

  private formatFileSize(bytes: number): string {
    const sizes = ['B', 'KB', 'MB', 'GB']
    if (bytes === 0) return '0 B'
    const i = Math.floor(Math.log(bytes) / Math.log(1024))
    return Math.round(bytes / Math.pow(1024, i) * 100) / 100 + ' ' + sizes[i]
  }
}

export const databaseRestore = new DatabaseRestore()

if (import.meta.url === `file://${process.argv[1]}`) {
  async function main() {
    const command = process.argv[2]
    const restoreTool = new DatabaseRestore()

    switch (command) {
      case 'restore':
        const backupFile = process.argv[3]
        await restoreTool.restoreDatabase({ backupFile })
        break
      case 'list':
        const backups = await restoreTool.listBackups()
        console.log('可用备份文件:')
        backups.forEach(file => console.log(`  - ${file}`))
        break
      case 'latest':
        const latest = await restoreTool.findLatestBackup()
        if (latest) {
          console.log(`最新备份文件: ${latest}`)
          const info = await restoreTool.getBackupInfo(latest)
          console.log(`文件大小: ${info.formattedSize}`)
          console.log(`修改时间: ${info.mtime}`)
        } else {
          console.log('未找到备份文件')
        }
        break
      default:
        console.log('用法: node restore.ts [restore <file>|list|latest]')
        console.log('  restore <file>  - 恢复指定备份文件')
        console.log('  list            - 列出所有备份文件')
        console.log('  latest          - 显示最新备份文件信息')
        process.exit(1)
    }
  }

  main().catch(console.error)
}