doc_ai_backed/cmd/migrate/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"time"

	"gitea.com/texpixel/document_ai/config"
	"gitea.com/texpixel/document_ai/internal/storage/dao"
	"github.com/spf13/viper"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func main() {
	// 解析命令行参数
	testEnv := flag.String("test-env", "dev", "测试环境配置 (dev/prod)")
	prodEnv := flag.String("prod-env", "prod", "生产环境配置 (dev/prod)")
	flag.Parse()

	// 加载测试环境配置
	testConfigPath := fmt.Sprintf("./config/config_%s.yaml", *testEnv)
	testConfig, err := loadDatabaseConfig(testConfigPath)
	if err != nil {
		log.Fatalf("加载测试环境配置失败: %v", err)
	}

	// 连接测试数据库
	testDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Asia%%2FShanghai",
		testConfig.Username, testConfig.Password, testConfig.Host, testConfig.Port, testConfig.DBName)
	testDB, err := gorm.Open(mysql.Open(testDSN), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	})
	if err != nil {
		log.Fatalf("连接测试数据库失败: %v", err)
	}

	// 加载生产环境配置
	prodConfigPath := fmt.Sprintf("./config/config_%s.yaml", *prodEnv)
	prodConfig, err := loadDatabaseConfig(prodConfigPath)
	if err != nil {
		log.Fatalf("加载生产环境配置失败: %v", err)
	}

	// 连接生产数据库
	prodDSN := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Asia%%2FShanghai",
		prodConfig.Username, prodConfig.Password, prodConfig.Host, prodConfig.Port, prodConfig.DBName)
	prodDB, err := gorm.Open(mysql.Open(prodDSN), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Info),
	})
	if err != nil {
		log.Fatalf("连接生产数据库失败: %v", err)
	}

	// 执行迁移
	if err := migrateData(testDB, prodDB); err != nil {
		log.Fatalf("数据迁移失败: %v", err)
	}

	log.Println("数据迁移完成！")
}

func migrateData(testDB, prodDB *gorm.DB) error {
	_ = context.Background() // 保留以备将来使用

	// 从测试数据库读取所有任务数据（包含结果）
	type TaskWithResult struct {
		// Task 字段
		TaskID        int64     `gorm:"column:id"`
		UserID        int64     `gorm:"column:user_id"`
		TaskUUID      string    `gorm:"column:task_uuid"`
		FileName      string    `gorm:"column:file_name"`
		FileHash      string    `gorm:"column:file_hash"`
		FileURL       string    `gorm:"column:file_url"`
		TaskType      string    `gorm:"column:task_type"`
		Status        int       `gorm:"column:status"`
		CompletedAt   time.Time `gorm:"column:completed_at"`
		Remark        string    `gorm:"column:remark"`
		IP            string    `gorm:"column:ip"`
		TaskCreatedAt time.Time `gorm:"column:created_at"`
		TaskUpdatedAt time.Time `gorm:"column:updated_at"`
		// Result 字段
		ResultID        *int64     `gorm:"column:result_id"`
		ResultTaskID    *int64     `gorm:"column:result_task_id"`
		ResultTaskType  *string    `gorm:"column:result_task_type"`
		Latex           *string    `gorm:"column:latex"`
		Markdown        *string    `gorm:"column:markdown"`
		MathML          *string    `gorm:"column:mathml"`
		ResultCreatedAt *time.Time `gorm:"column:result_created_at"`
		ResultUpdatedAt *time.Time `gorm:"column:result_updated_at"`
	}

	var tasksWithResults []TaskWithResult
	query := `
		SELECT 
			t.id,
			t.user_id,
			t.task_uuid,
			t.file_name,
			t.file_hash,
			t.file_url,
			t.task_type,
			t.status,
			t.completed_at,
			t.remark,
			t.ip,
			t.created_at,
			t.updated_at,
			r.id as result_id,
			r.task_id as result_task_id,
			r.task_type as result_task_type,
			r.latex,
			r.markdown,
			r.mathml,
			r.created_at as result_created_at,
			r.updated_at as result_updated_at
		FROM recognition_tasks t
		LEFT JOIN recognition_results r ON t.id = r.task_id
		ORDER BY t.id
	`

	if err := testDB.Raw(query).Scan(&tasksWithResults).Error; err != nil {
		return fmt.Errorf("读取测试数据失败: %v", err)
	}

	log.Printf("从测试数据库读取到 %d 条任务记录", len(tasksWithResults))

	successCount := 0
	skipCount := 0
	errorCount := 0

	// 为每个任务使用独立事务，确保单个任务失败不影响其他任务
	for i, item := range tasksWithResults {
		// 开始事务
		tx := prodDB.Begin()

		// 检查生产数据库中是否已存在相同的 task_uuid
		var existingTask dao.RecognitionTask
		err := tx.Where("task_uuid = ?", item.TaskUUID).First(&existingTask).Error
		if err == nil {
			log.Printf("[%d/%d] 跳过已存在的任务: task_uuid=%s, id=%d", i+1, len(tasksWithResults), item.TaskUUID, existingTask.ID)
			tx.Rollback()
			skipCount++
			continue
		}
		if err != gorm.ErrRecordNotFound {
			log.Printf("[%d/%d] 检查任务是否存在时出错: task_uuid=%s, error=%v", i+1, len(tasksWithResults), item.TaskUUID, err)
			tx.Rollback()
			errorCount++
			continue
		}

		// 创建新任务（不指定ID，让数据库自动生成）
		newTask := &dao.RecognitionTask{
			UserID:      item.UserID,
			TaskUUID:    item.TaskUUID,
			FileName:    item.FileName,
			FileHash:    item.FileHash,
			FileURL:     item.FileURL,
			TaskType:    dao.TaskType(item.TaskType),
			Status:      dao.TaskStatus(item.Status),
			CompletedAt: item.CompletedAt,
			Remark:      item.Remark,
			IP:          item.IP,
		}
		// 保留原始时间戳
		newTask.CreatedAt = item.TaskCreatedAt
		newTask.UpdatedAt = item.TaskUpdatedAt

		if err := tx.Create(newTask).Error; err != nil {
			log.Printf("[%d/%d] 创建任务失败: task_uuid=%s, error=%v", i+1, len(tasksWithResults), item.TaskUUID, err)
			tx.Rollback()
			errorCount++
			continue
		}

		log.Printf("[%d/%d] 创建任务成功: task_uuid=%s, 新ID=%d", i+1, len(tasksWithResults), item.TaskUUID, newTask.ID)

		// 如果有结果数据，创建结果记录
		if item.ResultID != nil {
			// 处理可能为NULL的字段
			latex := ""
			if item.Latex != nil {
				latex = *item.Latex
			}
			markdown := ""
			if item.Markdown != nil {
				markdown = *item.Markdown
			}
			mathml := ""
			if item.MathML != nil {
				mathml = *item.MathML
			}

			newResult := dao.RecognitionResult{
				TaskID:   newTask.ID, // 使用新任务的ID
				TaskType: dao.TaskType(item.TaskType),
				Latex:    latex,
				Markdown: markdown,
				MathML:   mathml,
			}
			// 保留原始时间戳
			if item.ResultCreatedAt != nil {
				newResult.CreatedAt = *item.ResultCreatedAt
			}
			if item.ResultUpdatedAt != nil {
				newResult.UpdatedAt = *item.ResultUpdatedAt
			}

			if err := tx.Create(&newResult).Error; err != nil {
				log.Printf("[%d/%d] 创建结果失败: task_id=%d, error=%v", i+1, len(tasksWithResults), newTask.ID, err)
				tx.Rollback() // 回滚整个事务（包括任务）
				errorCount++
				continue
			}

			log.Printf("[%d/%d] 创建结果成功: task_id=%d", i+1, len(tasksWithResults), newTask.ID)
		}

		// 提交事务
		if err := tx.Commit().Error; err != nil {
			log.Printf("[%d/%d] 提交事务失败: task_uuid=%s, error=%v", i+1, len(tasksWithResults), item.TaskUUID, err)
			errorCount++
			continue
		}

		successCount++
	}

	log.Printf("迁移完成统计:")
	log.Printf("  成功: %d 条", successCount)
	log.Printf("  跳过: %d 条", skipCount)
	log.Printf("  失败: %d 条", errorCount)

	return nil
}

// loadDatabaseConfig 从配置文件加载数据库配置
func loadDatabaseConfig(configPath string) (config.DatabaseConfig, error) {
	v := viper.New()
	v.SetConfigFile(configPath)
	if err := v.ReadInConfig(); err != nil {
		return config.DatabaseConfig{}, err
	}

	var dbConfig config.DatabaseConfig
	if err := v.UnmarshalKey("database", &dbConfig); err != nil {
		return config.DatabaseConfig{}, err
	}

	return dbConfig, nil
}