cleanup-postgresql-migrations.js

#!/usr/bin/env node

/**
 * PostgreSQL 迁移失败清理脚本
 * 用于清理 AnythingLLM 项目中的失败迁移记录
 */

const { Client } = require('pg');
const readline = require('readline');
const fs = require('fs');
const path = require('path');

// 创建命令行界面
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});

// 配置
const config = {
  host: 'rag-ai.e.aivencloud.com',
  port: 15005,
  ssl: {
    rejectUnauthorized: false // Aiven 通常需要这个设置
  }
};

// 辅助函数：提问
function question(query, hidden = false) {
  return new Promise((resolve) => {
    if (hidden) {
      const stdin = process.openStdin();
      process.stdout.write(query);
      stdin.setEncoding('utf8');
      stdin.on('data', (char) => {
        char = char.toString();
        if (char === '\n' || char === '\r' || char === '\u0004') {
          stdin.pause();
          resolve();
        } else {
          process.stdout.write('*');
        }
      });
    } else {
      rl.question(query, resolve);
    }
  });
}

// 主函数
async function main() {
  console.log('\n=== PostgreSQL 迁移失败清理工具 ===\n');
  
  try {
    // 获取数据库连接信息
    const database = await question('请输入数据库名 [默认: anythingllm]: ') || 'anythingllm';
    const username = await question('请输入用户名 [默认: postgres]: ') || 'postgres';
    const password = await question('请输入密码: ', true);
    
    config.user = username;
    config.password = password;
    config.database = database;
    
    // 创建客户端
    const client = new Client(config);
    
    console.log('\n正在连接到数据库...');
    await client.connect();
    console.log('✓ 连接成功\n');
    
    // 检查迁移表
    console.log('检查 _prisma_migrations 表状态...');
    const migrationsResult = await client.query(`
      SELECT 
        migration_name,
        started_at,
        finished_at,
        logs
      FROM _prisma_migrations 
      ORDER BY started_at DESC
    `);
    
    if (migrationsResult.rows.length === 0) {
      console.log('未找到迁移记录，数据库可能是干净的。');
      await client.end();
      rl.close();
      return;
    }
    
    // 显示迁移状态
    console.log('\n当前迁移状态:');
    console.log('─'.repeat(80));
    console.log('迁移名称'.padEnd(30) + '状态'.padEnd(15) + '开始时间');
    console.log('─'.repeat(80));
    
    migrationsResult.rows.forEach(row => {
      const status = row.finished_at ? '已完成' : '失败';
      const time = row.started_at.toLocaleString('zh-CN');
      console.log(`${row.migration_name.padEnd(30)}${status.padEnd(15)}${time}`);
    });
    
    // 查找失败的迁移
    const failedMigrations = migrationsResult.rows.filter(row => !row.finished_at);
    
    if (failedMigrations.length === 0) {
      console.log('\n✓ 没有找到失败的迁移记录。');
      await client.end();
      rl.close();
      return;
    }
    
    console.log(`\n发现 ${failedMigrations.length} 个失败的迁移:`);
    failedMigrations.forEach(row => {
      console.log(`- ${row.migration_name}`);
    });
    
    // 选择清理策略
    console.log('\n请选择清理策略:');
    console.log('1. 仅删除失败的迁移记录');
    console.log('2. 删除所有未完成的迁移');
    console.log('3. 完全重置迁移表（谨慎使用）');
    console.log('0. 退出');
    
    const choice = await question('\n请输入选择 [0-3]: ');
    
    switch (choice) {
      case '1':
        await cleanupFailedOnly(client, failedMigrations);
        break;
      case '2':
        await cleanupAllIncomplete(client);
        break;
      case '3':
        await resetMigrationsTable(client);
        break;
      case '0':
        console.log('退出程序。');
        break;
      default:
        console.log('无效选择。');
    }
    
    // 检查数据库表
    console.log('\n检查数据库表结构...');
    const tablesResult = await client.query(`
      SELECT table_name 
      FROM information_schema.tables 
      WHERE table_schema = 'public' 
      AND table_type = 'BASE TABLE'
      ORDER BY table_name
    `);
    
    console.log('\n当前数据库中的表:');
    tablesResult.rows.forEach(row => {
      console.log(`- ${row.table_name}`);
    });
    
    await client.end();
    
  } catch (error) {
    console.error('\n错误:', error.message);
    if (error.code === '3D000') {
      console.log('数据库不存在，请先创建数据库。');
    } else if (error.code === '28P01') {
      console.log('认证失败，请检查用户名和密码。');
    } else if (error.code === '08006') {
      console.log('连接失败，请检查网络和连接参数。');
    }
  }
  
  rl.close();
}

// 清理失败的迁移
async function cleanupFailedOnly(client, failedMigrations) {
  console.log('\n=== 清理失败的迁移 ===');
  
  // 备份
  const backup = await question('是否备份失败记录？[Y/n]: ');
  if (backup.toLowerCase() !== 'n') {
    await client.query(`
      CREATE TABLE _prisma_migrations_failed_backup AS 
      SELECT * FROM _prisma_migrations 
      WHERE finished_at IS NULL
    `);
    console.log('✓ 备份完成');
  }
  
  // 删除失败记录
  for (const migration of failedMigrations) {
    await client.query(`
      DELETE FROM _prisma_migrations 
      WHERE migration_name = '${migration.migration_name}'
    `);
    console.log(`✓ 已删除: ${migration.migration_name}`);
  }
  
  console.log('\n✓ 清理完成！');
  console.log('\n现在可以重新运行迁移命令:');
  console.log('npx prisma migrate dev');
}

// 清理所有未完成的迁移
async function cleanupAllIncomplete(client) {
  console.log('\n=== 清理所有未完成的迁移 ===');
  
  const confirm = await question('这将删除所有未完成的迁移记录，确定吗？[y/N]: ');
  if (confirm.toLowerCase() !== 'y') {
    console.log('操作已取消。');
    return;
  }
  
  // 备份
  await client.query(`
    CREATE TABLE _prisma_migrations_incomplete_backup AS 
    SELECT * FROM _prisma_migrations 
    WHERE finished_at IS NULL
  `);
  console.log('✓ 备份完成');
  
  // 删除
  const result = await client.query(`
    DELETE FROM _prisma_migrations 
    WHERE finished_at IS NULL
  `);
  
  console.log(`✓ 已删除 ${result.rowCount} 个未完成的迁移`);
  console.log('\n✓ 清理完成！');
}

// 重置迁移表
async function resetMigrationsTable(client) {
  console.log('\n=== 重置迁移表 ===');
  
  const confirm1 = await question('这将完全清空迁移表，确定吗？[y/N]: ');
  if (confirm1.toLowerCase() !== 'y') {
    console.log('操作已取消。');
    return;
  }
  
  const confirm2 = await question('这可能导致数据丢失，再次确认？[y/N]: ');
  if (confirm2.toLowerCase() !== 'y') {
    console.log('操作已取消。');
    return;
  }
  
  // 备份
  await client.query(`
    CREATE TABLE _prisma_migrations_full_backup AS 
    SELECT * FROM _prisma_migrations
  `);
  console.log('✓ 完整备份完成');
  
  // 清空表
  await client.query('TRUNCATE _prisma_migrations');
  console.log('✓ 迁移表已重置');
  
  console.log('\n⚠️  警告：迁移表已完全重置');
  console.log('下次运行迁移时将从头开始创建所有表结构');
}

// 运行主函数
if (require.main === module) {
  main().catch(console.error);
}

module.exports = {
  main,
  cleanupFailedOnly,
  cleanupAllIncomplete,
  resetMigrationsTable
};