observe-server/cmd/migrate/server.go

package migrate

import (
	"bytes"
	"fmt"
	"strconv"
	"text/template"
	"time"

	"github.com/go-admin-team/go-admin-core/sdk"
	"github.com/go-admin-team/go-admin-core/sdk/pkg"
	"github.com/pkg/errors"

	"github.com/go-admin-team/go-admin-core/config/source/file"
	"github.com/spf13/cobra"

	"go-admin/cmd/migrate/migration"
	_ "go-admin/cmd/migrate/migration/version"

	extConfig "go-admin/config"

	"github.com/go-admin-team/go-admin-core/sdk/config"

	_ "go-admin/cmd/migrate/migration/version-local"
	"go-admin/common/database"
	"go-admin/common/models"
	"go-admin/common/olap"
)

var (
	configYml string
	generate  bool
	goAdmin   bool
	host      string
	StartCmd  = &cobra.Command{
		Use:     "migrate",
		Short:   "Initialize the database",
		Example: "go-admin migrate -c config/settings.yml",
		PreRun: func(cmd *cobra.Command, args []string) {
			// 注入配置扩展项, 后面的clickhouse迁移需要这些config
			config.ExtendConfig = &extConfig.ExtConfig
		},
		Run: func(cmd *cobra.Command, args []string) {
			run()
		},
	}
)

// fixme 在您看不见代码的时候运行迁移，我觉得是不安全的，所以编译后最好不要去执行迁移
func init() {
	StartCmd.PersistentFlags().StringVarP(&configYml, "config", "c", "config/settings.yml", "Start server with provided configuration file")
	StartCmd.PersistentFlags().BoolVarP(&generate, "generate", "g", false, "generate migration file")
	StartCmd.PersistentFlags().BoolVarP(&goAdmin, "goAdmin", "a", false, "generate go-admin migration file")
	StartCmd.PersistentFlags().StringVarP(&host, "domain", "d", "*", "select tenant host")
}

func run() {

	if !generate {
		fmt.Println(`start init`)
		//1. 读取配置
		config.Setup(
			file.NewSource(file.WithPath(configYml)),
			initDB,
		)
	} else {
		fmt.Println(`generate migration file`)
		_ = genFile()
	}
}

func migrateModel() error {
	if host == "" {
		host = "*"
	}
	db := sdk.Runtime.GetDbByKey(host)
	if config.DatabasesConfig[host].Driver == "mysql" {
		//初始化数据库时候用
		db.Set("gorm:table_options", "ENGINE=InnoDB CHARSET=utf8mb4")
	}
	err := db.Debug().AutoMigrate(&models.Migration{})
	if err != nil {
		return err
	}
	migration.Migrate.SetDb(db.Debug())
	migration.Migrate.Migrate()
	return err
}

func migrateClickhouseModel() error {
	db := olap.GetClickhouseOrm()
	if err := db.Exec(`CREATE TABLE if not exists otel.sys_migration_local on cluster default (
		version String,
		apply_time DateTime64(3)
	) ENGINE = MergeTree ORDER BY tuple() SETTINGS index_granularity = 8192`).Error; err != nil {
		return errors.Wrap(err, "clickhouse生成sys_migration_local 失败")
	}
	if err := db.Exec(`create table if not exists otel.sys_migration on cluster default(
		version String,
		apply_time DateTime64(3)
	) engine = Distributed('default', 'otel', 'sys_migration_local', cityHash64(version))`).Error; err != nil {
		return errors.Wrap(err, "clickhouse生成sys_migration分布式表失败")
	}
	err := db.Debug().AutoMigrate(&models.Migration{})
	if err != nil {
		return errors.Wrap(err, "auto migrate faild")
	}
	migration.Migrate.SetDb(db.Debug())
	migration.Migrate.Migrate()
	return err
}

func initDB() {
	//3. 初始化数据库链接
	database.Setup()
	// olap.Setup()
	//4. 数据库迁移
	fmt.Println("数据库迁移开始")
	_ = migrateModel()
	fmt.Println(`数据库基础数据初始化成功`)
	// fmt.Printf("数据库clickhouse迁移开始")
	// if err := migrateClickhouseModel(); err != nil {
	// 	panic(err)
	// }
	// fmt.Printf("数据库clickhouse迁移成功")
}

func genFile() error {
	t1, err := template.ParseFiles("template/migrate.template")
	if err != nil {
		return err
	}
	m := map[string]string{}
	m["GenerateTime"] = strconv.FormatInt(time.Now().UnixNano()/1e6, 10)
	m["Package"] = "version_local"
	if goAdmin {
		m["Package"] = "version"
	}
	var b1 bytes.Buffer
	if err = t1.Execute(&b1, m); err != nil {
		return err
	}
	if goAdmin {
		pkg.FileCreate(b1, "./cmd/migrate/migration/version/"+m["GenerateTime"]+"_migrate.go")
	} else {
		pkg.FileCreate(b1, "./cmd/migrate/migration/version-local/"+m["GenerateTime"]+"_migrate.go")
	}
	return nil
}