[优化] go 引用调整

2026-10-03 07:06:36 +08:00 · 2026-06-06 10:26:20 +08:00
parent ee1110b752
commit 3cfefb4367
552 changed files with 1642 additions and 2185 deletions
@@ -0,0 +1,369 @@
+package model
+
+import (
+	"fmt"
+	"log/slog"
+	"os"
+	"reflect"
+	"strings"
+	"sync"
+
+	"github.com/rain-kl/openflare/openflare-server/common"
+	"github.com/rain-kl/openflare/openflare-server/utils/security"
+
+	"github.com/glebarez/sqlite"
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+	"gorm.io/gorm/schema"
+)
+
+var DB *gorm.DB
+
+type dbModel struct {
+	value     any
+	tableName string
+	hasIDPK   bool
+}
+
+func registeredModels() []any {
+	return []any{
+		&User{},
+		&AuthSource{},
+		&ExternalAccount{},
+		&Option{},
+		&Origin{},
+		&ProxyRoute{},
+		&PagesProject{},
+		&PagesDeployment{},
+		&PagesDeploymentFile{},
+		&ConfigVersion{},
+		&Node{},
+
+		&NodeSystemProfile{},
+		&ApplyLog{},
+		&NodeMetricSnapshot{},
+		&NodeRequestReport{},
+		&NodeAccessLog{},
+		&NodeHealthEvent{},
+		&NodeObservationOpenresty{},
+		&NodeObservationFrps{},
+		&NodeObservationFrpc{},
+		&TLSCertificate{},
+		&ManagedDomain{},
+		&AcmeAccount{},
+		&DnsAccount{},
+		&WAFRuleGroup{},
+		&WAFIPGroup{},
+		&WAFRuleGroupBinding{},
+	}
+}
+
+func currentSchemaMetadataModels() []any {
+	return nil
+}
+
+func legacySchemaMetadataModels() []any {
+	return []any{
+		&DatabaseSchemaVersion{},
+	}
+}
+
+func schemaMetadataModels() []any {
+	models := make([]any, 0, len(currentSchemaMetadataModels())+len(legacySchemaMetadataModels()))
+	models = append(models, currentSchemaMetadataModels()...)
+	models = append(models, legacySchemaMetadataModels()...)
+	return models
+}
+
+func buildDBModels() ([]dbModel, error) {
+	models := registeredModels()
+	result := make([]dbModel, 0, len(models))
+	namer := schema.NamingStrategy{}
+	cache := &sync.Map{}
+	for _, item := range models {
+		parsed, err := schema.Parse(item, cache, namer)
+		if err != nil {
+			return nil, err
+		}
+		hasIDPK := len(parsed.PrimaryFields) == 1 && parsed.PrimaryFields[0].DBName == "id"
+		result = append(result, dbModel{
+			value:     item,
+			tableName: parsed.Table,
+			hasIDPK:   hasIDPK,
+		})
+	}
+	return result, nil
+}
+
+func createRootAccountIfNeed() error {
+	var user User
+	//if user.Status != common.UserStatusEnabled {
+	if err := DB.First(&user).Error; err != nil {
+		slog.Info("no user exists, create a root user", "username", "root")
+		hashedPassword, err := security.Password2Hash("123456")
+		if err != nil {
+			return err
+		}
+		rootUser := User{
+			Username:    "root",
+			Password:    hashedPassword,
+			Role:        common.RoleRootUser,
+			Status:      common.UserStatusEnabled,
+			DisplayName: "Root User",
+		}
+		DB.Create(&rootUser)
+	}
+	return nil
+}
+
+func CountTable(tableName string) (num int64) {
+	DB.Table(tableName).Count(&num)
+	return
+}
+
+func openDatabase() (*gorm.DB, string, error) {
+	if common.SQLDSN != "" {
+		db, err := gorm.Open(postgres.Open(common.SQLDSN), &gorm.Config{})
+		if err != nil {
+			return nil, "", err
+		}
+		return db, "postgres", nil
+	}
+	db, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{})
+	if err != nil {
+		return nil, "", err
+	}
+	slog.Info("database DSN not set, using SQLite as database", "sqlite_path", common.SQLitePath)
+	return db, "sqlite", nil
+}
+
+func autoMigrateAll(db *gorm.DB) error {
+	return autoMigrateAllExcept(db, nil)
+}
+
+func autoMigrateAllExcept(db *gorm.DB, excludedTables map[string]bool) error {
+	models := registeredModels()
+	for i, item := range models {
+		name := fmt.Sprintf("%T", item)
+		tableName, err := tableNameForModel(item)
+		if err != nil {
+			return fmt.Errorf("resolve table name for %s failed: %w", name, err)
+		}
+		if excludedTables[tableName] {
+			slog.Info("autoMigrateAll: skipped model", "index", fmt.Sprintf("%d/%d", i+1, len(models)), "model", name, "table", tableName)
+			continue
+		}
+		slog.Info("autoMigrateAll: migrating model", "index", fmt.Sprintf("%d/%d", i+1, len(models)), "model", name)
+		if err := db.AutoMigrate(item); err != nil {
+			return fmt.Errorf("AutoMigrate %s failed: %w", name, err)
+		}
+		slog.Info("autoMigrateAll: migrated model", "model", name)
+	}
+	return nil
+}
+
+func tableNameForModel(item any) (string, error) {
+	namer := schema.NamingStrategy{}
+	cache := &sync.Map{}
+	parsed, err := schema.Parse(item, cache, namer)
+	if err != nil {
+		return "", err
+	}
+	return parsed.Table, nil
+}
+
+func isDatabaseEmpty(db *gorm.DB) (bool, error) {
+	models, err := buildDBModels()
+	if err != nil {
+		return false, err
+	}
+	for _, item := range models {
+		if isShardedObservabilityTable(item.tableName) {
+			for _, table := range observabilityShardTables(item.tableName) {
+				if !db.Migrator().HasTable(table) {
+					continue
+				}
+				var count int64
+				if err := db.Table(table).Limit(1).Count(&count).Error; err != nil {
+					return false, err
+				}
+				if count > 0 {
+					return false, nil
+				}
+			}
+			continue
+		}
+		if !db.Migrator().HasTable(item.value) {
+			continue
+		}
+		var count int64
+		if err := db.Model(item.value).Limit(1).Count(&count).Error; err != nil {
+			return false, err
+		}
+		if count > 0 {
+			return false, nil
+		}
+	}
+	return true, nil
+}
+
+func sqliteSourceExists() bool {
+	info, err := os.Stat(common.SQLitePath)
+	if err != nil {
+		return false
+	}
+	return !info.IsDir()
+}
+
+func migrateSQLiteDataIfNeeded(target *gorm.DB, backend string) error {
+	if backend != "postgres" {
+		return nil
+	}
+	empty, err := isDatabaseEmpty(target)
+	if err != nil {
+		return err
+	}
+	if !empty {
+		slog.Info("skip sqlite migration because target database already has data", "backend", backend)
+		return nil
+	}
+	if !sqliteSourceExists() {
+		slog.Info("skip sqlite migration because sqlite source file was not found", "sqlite_path", common.SQLitePath)
+		return nil
+	}
+
+	source, err := gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{
+		PrepareStmt: true,
+	})
+	if err != nil {
+		return fmt.Errorf("open sqlite source database failed: %w", err)
+	}
+	sourceSQLDB, err := source.DB()
+	if err != nil {
+		return fmt.Errorf("get sqlite source database handle failed: %w", err)
+	}
+	defer func() {
+		_ = sourceSQLDB.Close()
+	}()
+
+	models, err := buildDBModels()
+	if err != nil {
+		return err
+	}
+
+	slog.Info("starting sqlite to postgres database migration", "sqlite_path", common.SQLitePath)
+	err = target.Transaction(func(tx *gorm.DB) error {
+		for _, item := range models {
+			if err := migrateTableData(source, tx, item); err != nil {
+				return err
+			}
+			if item.hasIDPK {
+				if err := resetPostgresSequence(tx, item.tableName); err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	slog.Info("sqlite to postgres database migration completed", "sqlite_path", common.SQLitePath)
+	return nil
+}
+
+func migrateTableData(source *gorm.DB, target *gorm.DB, item dbModel) error {
+	if !source.Migrator().HasTable(item.value) {
+		slog.Info("database migration progress", "table", item.tableName, "migrated", 0, "total", 0, "status", "skipped_missing_source_table")
+		return nil
+	}
+	var total int64
+	if err := source.Model(item.value).Count(&total).Error; err != nil {
+		return fmt.Errorf("count sqlite table %s failed: %w", item.tableName, err)
+	}
+	slog.Info("database migration progress", "table", item.tableName, "migrated", 0, "total", total, "status", "starting")
+	if total == 0 {
+		slog.Info("database migration progress", "table", item.tableName, "migrated", 0, "total", total, "status", "completed")
+		return nil
+	}
+
+	modelType := reflect.TypeOf(item.value).Elem()
+	sliceType := reflect.SliceOf(modelType)
+	migrated := int64(0)
+	offset := 0
+	const batchSize = 200
+
+	for {
+		batchPtr := reflect.New(sliceType)
+		query := source.Model(item.value).Limit(batchSize).Offset(offset)
+		if item.hasIDPK {
+			query = query.Order("id ASC")
+		}
+		if err := query.Find(batchPtr.Interface()).Error; err != nil {
+			return fmt.Errorf("read sqlite table %s failed: %w", item.tableName, err)
+		}
+		batchLen := batchPtr.Elem().Len()
+		if batchLen == 0 {
+			break
+		}
+		if isShardedObservabilityTable(item.tableName) {
+			for index := 0; index < batchLen; index++ {
+				record := batchPtr.Elem().Index(index)
+				if err := target.Create(record.Addr().Interface()).Error; err != nil {
+					return fmt.Errorf("write target sharded table %s failed: %w", item.tableName, err)
+				}
+			}
+		} else {
+			if err := target.Create(batchPtr.Interface()).Error; err != nil {
+				return fmt.Errorf("write target table %s failed: %w", item.tableName, err)
+			}
+		}
+		migrated += int64(batchLen)
+		offset += batchLen
+		slog.Info("database migration progress", "table", item.tableName, "migrated", migrated, "total", total, "status", "running")
+	}
+
+	slog.Info("database migration progress", "table", item.tableName, "migrated", migrated, "total", total, "status", "completed")
+	return nil
+}
+
+func resetPostgresSequence(db *gorm.DB, tableName string) error {
+	sql := fmt.Sprintf(
+		"SELECT setval(pg_get_serial_sequence('%s', 'id'), COALESCE(MAX(id), 1), MAX(id) IS NOT NULL) FROM \"%s\"",
+		tableName,
+		tableName,
+	)
+	return db.Exec(sql).Error
+}
+
+func InitDB() (err error) {
+	db, backend, err := openDatabase()
+	if err != nil {
+		slog.Error("open database failed", "error", err)
+		os.Exit(1)
+	}
+	DB = db
+	if err = registerSharding(db, backend); err != nil {
+		return err
+	}
+	if err = ensureDatabaseSchemaUpToDate(db, backend); err != nil {
+		return err
+	}
+	return createRootAccountIfNeed()
+}
+
+func CloseDB() error {
+	sqlDB, err := DB.DB()
+	if err != nil {
+		return err
+	}
+	err = sqlDB.Close()
+	return err
+}
+
+func IsUniqueConstraintError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return strings.Contains(strings.ToLower(err.Error()), "unique")
+}