Init

internal/migration/migration.go

@@ -0,0 +1,101 @@
+package migration
+
+import (
+	"cs-agent/internal/models"
+	"cs-agent/internal/services"
+	"errors"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/mlogclub/simple/sqls"
+	"github.com/spf13/cast"
+)
+
+var migrationFuncs = make(map[int64]MigrationFunc)
+var versions = make([]int64, 0)
+var migrations = make(map[int64]models.Migration, 0)
+var mu sync.Mutex
+
+type MigrationFunc struct {
+	Version int64
+	Remark  string
+	Fn      func() error
+}
+
+func Migrate() error {
+	mu.Lock()
+	defer mu.Unlock()
+
+	if list := services.MigrationService.Find(sqls.NewCnd().Asc("version")); len(list) > 0 {
+		for _, element := range list {
+			migrations[element.Version] = element
+		}
+	}
+
+	for _, version := range versions {
+		if err := runMigration(version); err != nil {
+			slog.Error("migrate failed", "version", version, "error", err)
+			return err
+		}
+	}
+	return nil
+}
+
+func register(version int64, remark string, fn func() error) {
+	if len(versions) == 0 || version > versions[len(versions)-1] {
+		versions = append(versions, version)
+		migrationFuncs[version] = MigrationFunc{
+			Version: version,
+			Remark:  remark,
+			Fn:      fn,
+		}
+	} else {
+		slog.Error("register migration failed, version is less than latest version", slog.Any("version", version))
+		panic(errors.New("register migration failed, version is less than latest version. version: " + cast.ToString(version)))
+	}
+}
+
+func runMigration(version int64) error {
+	migration, found := migrations[version]
+	if found && migration.Success {
+		return nil
+	}
+
+	f, ok := migrationFuncs[version]
+	if !ok {
+		return errors.New("migration function not found")
+	}
+
+	err := f.Fn()
+
+	if !found {
+		migration = models.Migration{
+			Version:    f.Version,
+			Remark:     f.Remark,
+			Success:    false,
+			RetryCount: 0,
+			CreatedAt:  time.Now(),
+			UpdatedAt:  time.Now(),
+		}
+	}
+	if err == nil {
+		migration.Success = true
+	} else {
+		migration.Success = false
+		migration.ErrorInfo = err.Error()
+	}
+	migration.RetryCount++
+	migration.UpdatedAt = time.Now()
+	if found {
+		if e := services.MigrationService.Update(&migration); e != nil {
+			slog.Error("update migration failed", "version", version, "error", err)
+		}
+	} else {
+		if e := services.MigrationService.Create(&migration); e != nil {
+			slog.Error("create migration failed", "version", version, "error", err)
+		}
+	}
+
+	return err
+}