JCS-pub/coordinator/internal/cmd/migrate.go

package cmd

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/config"
	cortypes "gitlink.org.cn/cloudream/jcs-pub/coordinator/types"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

func init() {
	var configPath string
	cmd := cobra.Command{
		Use:   "migrate",
		Short: "Run database migrations",
		Run: func(cmd *cobra.Command, args []string) {
			migrate(configPath)
		},
	}
	cmd.Flags().StringVarP(&configPath, "config", "c", "", "Path to config file")
	RootCmd.AddCommand(&cmd)
}

func migrate(configPath string) {
	err := config.Init(configPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	db, err := gorm.Open(mysql.Open(config.Cfg().DB.MakeSourceString()))
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	db = db.Set("gorm:table_options", "CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci")

	migrateOne(db, cortypes.HubConnectivity{})
	migrateOne(db, cortypes.Hub{})
	migrateOne(db, cortypes.HubLocation{})
	migrateOne(db, cortypes.User{})
	migrateOne(db, cortypes.UserAccessToken{})
	migrateOne(db, cortypes.LoadedAccessToken{})

	fmt.Println("migrate success")
}

func migrateOne[T any](db *gorm.DB, model T) {
	err := db.AutoMigrate(model)
	if err != nil {
		fmt.Printf("migratting model %T: %v\n", model, err)
		os.Exit(1)
	}
}