go-skeleton/database/migration.go

package database

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"time"

	"go.uber.org/zap"
)

// MigrationConfig 迁移配置
type MigrationConfig struct {
	Environment string // 例如 local_postgres、production_mysql
	Directory   string // 按数据库方言隔离的迁移目录
	Timeout     int    // 超时时间（秒）
}

// MigrationManager 迁移管理器
type MigrationManager struct {
	config *MigrationConfig
	logger *zap.Logger
}

// NewMigrationManager 创建迁移管理器
func NewMigrationManager(config *MigrationConfig, logger *zap.Logger) *MigrationManager {
	return &MigrationManager{
		config: config,
		logger: logger,
	}
}

// CheckMigrations 检查是否有未应用的迁移
func (m *MigrationManager) CheckMigrations() error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(m.config.Timeout)*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "atlas", "migrate", "status", "--env", m.config.Environment)
	output, err := cmd.CombinedOutput()

	if err != nil {
		m.logger.Error("检查迁移状态失败",
			zap.Error(err),
			zap.String("output", string(output)))
		return fmt.Errorf("检查迁移状态失败: %w", err)
	}

	m.logger.Info("迁移状态检查完成", zap.String("output", string(output)))
	return nil
}

// GenerateMigration 生成迁移文件
func (m *MigrationManager) GenerateMigration(name string) error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(m.config.Timeout)*time.Second)
	defer cancel()

	var cmd *exec.Cmd
	if name != "" {
		cmd = exec.CommandContext(ctx, "atlas", "migrate", "diff", name, "--env", m.config.Environment)
	} else {
		cmd = exec.CommandContext(ctx, "atlas", "migrate", "diff", "--env", m.config.Environment)
	}

	output, err := cmd.CombinedOutput()

	if err != nil {
		m.logger.Error("生成迁移失败",
			zap.Error(err),
			zap.String("output", string(output)))
		return fmt.Errorf("生成迁移失败: %w", err)
	}

	m.logger.Info("迁移文件生成成功",
		zap.String("name", name),
		zap.String("output", string(output)))
	return nil
}

// ApplyMigrations 应用迁移
func (m *MigrationManager) ApplyMigrations(dryRun bool) error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(m.config.Timeout)*time.Second)
	defer cancel()

	args := []string{"migrate", "apply", "--env", m.config.Environment}
	if dryRun {
		args = append(args, "--dry-run")
	}
	cmd := exec.CommandContext(ctx, "atlas", args...)
	output, err := cmd.CombinedOutput()

	if err != nil {
		m.logger.Error("应用迁移失败",
			zap.Error(err),
			zap.String("output", string(output)))
		return fmt.Errorf("应用迁移失败: %w", err)
	}

	m.logger.Info("迁移应用成功",
		zap.Bool("dry_run", dryRun),
		zap.String("output", string(output)))
	return nil
}

// ValidateMigrations 验证迁移
func (m *MigrationManager) ValidateMigrations() error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(m.config.Timeout)*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "atlas", "migrate", "validate", "--env", m.config.Environment)
	output, err := cmd.CombinedOutput()

	if err != nil {
		m.logger.Error("验证迁移失败",
			zap.Error(err),
			zap.String("output", string(output)))
		return fmt.Errorf("验证迁移失败: %w", err)
	}

	m.logger.Info("迁移验证成功", zap.String("output", string(output)))
	return nil
}

// EnsureAtlasInstalled 确保Atlas已安装
func EnsureAtlasInstalled() error {
	cmd := exec.Command("atlas", "version")
	err := cmd.Run()
	if err != nil {
		return fmt.Errorf("Atlas CLI 未安装或不在 PATH 中，请安装: %w", err)
	}
	return nil
}

// InitMigrationDirectory 初始化迁移目录
func (m *MigrationManager) InitMigrationDirectory() error {
	migrationDir := m.config.Directory
	if migrationDir == "" {
		migrationDir = "migrations"
	}

	// 检查目录是否存在
	if _, err := os.Stat(migrationDir); os.IsNotExist(err) {
		if err := os.MkdirAll(migrationDir, 0755); err != nil {
			m.logger.Error("创建迁移目录失败", zap.Error(err))
			return fmt.Errorf("创建迁移目录失败: %w", err)
		}
		m.logger.Info("迁移目录创建成功", zap.String("dir", migrationDir))
	}

	return nil
}