Godemo/iot-base-station/cmd/migrate/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"

	"iot-base-station/internal/config"
	"iot-base-station/pkg/database"

	"github.com/golang-migrate/migrate/v4"
	_ "github.com/golang-migrate/migrate/v4/database/postgres"
	_ "github.com/golang-migrate/migrate/v4/source/file"
)

func main() {
	var (
		action = flag.String("action", "", "Migration action: up, down, version, force")
		step   = flag.Int("step", 0, "Number of steps to migrate (for up/down)")
		version = flag.Int("version", 0, "Target version (for force)")
	)
	flag.Parse()

	if *action == "" {
		fmt.Println("Usage: migrate -action=<up|down|version|force> [-step=N] [-version=N]")
		os.Exit(1)
	}

	// 加载配置
	cfg, err := config.LoadConfig("config/config.yaml")
	if err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}

	// 连接数据库
	db, err := database.InitPostgres(cfg.Database.Postgres)
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}
	defer db.Close()

	// 创建迁移实例
	m, err := migrate.New(
		"file://migrations",
		fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable",
			cfg.Database.Postgres.User,
			cfg.Database.Postgres.Password,
			cfg.Database.Postgres.Host,
			cfg.Database.Postgres.Port,
			cfg.Database.Postgres.DBName,
		),
	)
	if err != nil {
		log.Fatalf("Failed to create migration instance: %v", err)
	}
	defer m.Close()

	// 执行迁移
	switch *action {
	case "up":
		if *step > 0 {
			if err := m.Steps(*step); err != nil {
				log.Fatalf("Failed to migrate up: %v", err)
			}
		} else {
			if err := m.Up(); err != nil && err != migrate.ErrNoChange {
				log.Fatalf("Failed to migrate up: %v", err)
			}
		}
		fmt.Println("Migration up completed successfully")

	case "down":
		if *step > 0 {
			if err := m.Steps(-*step); err != nil {
				log.Fatalf("Failed to migrate down: %v", err)
			}
		} else {
			if err := m.Down(); err != nil && err != migrate.ErrNoChange {
				log.Fatalf("Failed to migrate down: %v", err)
			}
		}
		fmt.Println("Migration down completed successfully")

	case "version":
		version, dirty, err := m.Version()
		if err != nil {
			log.Fatalf("Failed to get version: %v", err)
		}
		fmt.Printf("Current version: %d, dirty: %v\n", version, dirty)

	case "force":
		if *version < 0 {
			log.Fatal("Version must be specified for force action")
		}
		if err := m.Force(*version); err != nil {
			log.Fatalf("Failed to force version: %v", err)
		}
		fmt.Printf("Forced version to %d\n", *version)

	default:
		log.Fatalf("Unknown action: %s", *action)
	}
}