Jelajahi Sumber

fix(database): make cross-db migration lossless, transactional, and pre-checked

migrationModels was missing ClientGroup and ClientGlobalTraffic, so both
migration directions silently dropped client groups and global client
traffic; the model list is now extracted to allModels and a parity test
keeps the two lists from drifting again. MigrateData runs its truncate
and copy inside one transaction so a failed import rolls back instead of
leaving the destination truncated (sequences resync after commit since
setval is non-transactional). New PrepareSQLiteForMigration rejects
uploads that are not a panel database and AutoMigrates old backups so
their missing tables cannot break the row copy.
MHSanaei 20 jam lalu
induk
melakukan
54fc0fd47c

+ 10 - 6
internal/database/db.go

@@ -56,8 +56,8 @@ const (
 	defaultPassword = "admin"
 )
 
-func initModels() error {
-	models := []any{
+func allModels() []any {
+	return []any{
 		&model.User{},
 		&model.Inbound{},
 		&model.OutboundTraffics{},
@@ -78,12 +78,16 @@ func initModels() error {
 		&model.ClientGlobalTraffic{},
 		&model.OutboundSubscription{},
 	}
+}
+
+func initModels() error {
+	models := allModels()
 	for _, mdl := range models {
 		if IsPostgres() && postgresModelSettled(mdl) {
 			continue
 		}
 		if err := db.AutoMigrate(mdl); err != nil {
-			if isIgnorableDuplicateColumnErr(err, mdl) {
+			if isIgnorableDuplicateColumnErr(db, err, mdl) {
 				log.Printf("Ignoring duplicate column during auto migration for %T: %v", mdl, err)
 				continue
 			}
@@ -999,7 +1003,7 @@ func dedupeInboundSettingsClients() error {
 	return nil
 }
 
-func isIgnorableDuplicateColumnErr(err error, mdl any) bool {
+func isIgnorableDuplicateColumnErr(gdb *gorm.DB, err error, mdl any) bool {
 	if err == nil {
 		return false
 	}
@@ -1009,14 +1013,14 @@ func isIgnorableDuplicateColumnErr(err error, mdl any) bool {
 	if _, after, ok := strings.Cut(errMsg, sqlitePrefix); ok {
 		col := strings.TrimSpace(after)
 		col = strings.Trim(col, "`\"[]")
-		return col != "" && db != nil && db.Migrator().HasColumn(mdl, col)
+		return col != "" && gdb != nil && gdb.Migrator().HasColumn(mdl, col)
 	}
 	if strings.Contains(errMsg, "already exists") && strings.Contains(errMsg, "column ") {
 		if _, after, ok := strings.Cut(errMsg, "column \""); ok {
 			rest := after
 			if e := strings.Index(rest, "\""); e > 0 {
 				col := rest[:e]
-				return col != "" && db != nil && db.Migrator().HasColumn(mdl, col)
+				return col != "" && gdb != nil && gdb.Migrator().HasColumn(mdl, col)
 			}
 		}
 	}

+ 66 - 24
internal/database/migrate_data.go

@@ -25,7 +25,8 @@ import (
 // related tests.
 //
 // Important: When adding a new top-level model (like OutboundSubscription),
-// you must add it here **in addition to** the list in internal/database/db.go:initModels().
+// you must add it here **in addition to** allModels() in internal/database/db.go;
+// TestMigrationModelsMatchPanelModels fails when the two lists drift apart.
 // This list is used for:
 //   - Creating the destination schema during cross-DB migration
 //   - Truncating tables
@@ -48,17 +49,21 @@ func migrationModels() []any {
 		&model.ClientRecord{},
 		&model.ClientInbound{},
 		&model.ClientExternalLink{},
+		&model.ClientGroup{},
 		&model.InboundFallback{},
 		&model.Host{},
 		&model.NodeClientTraffic{},
 		&model.NodeClientIp{},
+		&model.ClientGlobalTraffic{},
 		&model.OutboundSubscription{},
 	}
 }
 
 // MigrateData copies every row from the configured SQLite file at srcPath into
 // a fresh PostgreSQL database described by dstDSN. The destination tables are
-// (re)created with AutoMigrate before the copy. Source data is left untouched.
+// (re)created with AutoMigrate; truncate and copy then run in one transaction,
+// so a failed migration leaves the destination data unchanged. Source data is
+// left untouched.
 func MigrateData(srcPath, dstDSN string) error {
 	if _, err := os.Stat(srcPath); err != nil {
 		return fmt.Errorf("source sqlite not found at %s: %w", srcPath, err)
@@ -100,33 +105,41 @@ func MigrateData(srcPath, dstDSN string) error {
 		}
 	}
 
-	// AutoMigrate re-creates the legacy client_traffics -> inbounds foreign key,
-	// but the running panel drops it (see dropLegacyForeignKeys) and tolerates
-	// client_traffics rows whose inbound was deleted. Drop it here too so copying
-	// such orphaned rows can't fail with an fk_inbounds_client_stats violation.
-	if err := dst.Exec("ALTER TABLE client_traffics DROP CONSTRAINT IF EXISTS fk_inbounds_client_stats").Error; err != nil {
-		return fmt.Errorf("drop legacy foreign key: %w", err)
-	}
+	totalRows := 0
+	txErr := dst.Transaction(func(tx *gorm.DB) error {
+		// AutoMigrate re-creates the legacy client_traffics -> inbounds foreign key,
+		// but the running panel drops it (see dropLegacyForeignKeys) and tolerates
+		// client_traffics rows whose inbound was deleted. Drop it here too so copying
+		// such orphaned rows can't fail with an fk_inbounds_client_stats violation.
+		if err := tx.Exec("ALTER TABLE client_traffics DROP CONSTRAINT IF EXISTS fk_inbounds_client_stats").Error; err != nil {
+			return fmt.Errorf("drop legacy foreign key: %w", err)
+		}
 
-	// Empty the destination tables so the migration is idempotent: a fresh
-	// PostgreSQL DB already holds an auto-seeded admin (id=1) from any prior
-	// panel start, and a partially-failed earlier run leaves rows behind. Either
-	// way a plain INSERT with explicit ids would collide on users_pkey, so clear
-	// our tables (only) before copying.
-	if err := truncatePostgresTables(dst, migrationModels()); err != nil {
-		return fmt.Errorf("clear destination tables: %w", err)
-	}
+		// Empty the destination tables before copying: a fresh PostgreSQL DB
+		// already holds an auto-seeded admin (id=1) from any prior panel start,
+		// so a plain INSERT with explicit ids would collide on users_pkey. Only
+		// the panel's own tables are cleared, and a failure anywhere in this
+		// transaction rolls the clear back with everything else.
+		if err := truncatePostgresTables(tx, migrationModels()); err != nil {
+			return fmt.Errorf("clear destination tables: %w", err)
+		}
 
-	totalRows := 0
-	for _, m := range migrationModels() {
-		n, err := copyTable(src, dst, m)
-		if err != nil {
-			return fmt.Errorf("copy %T: %w", m, err)
+		for _, m := range migrationModels() {
+			n, err := copyTable(src, tx, m)
+			if err != nil {
+				return fmt.Errorf("copy %T: %w", m, err)
+			}
+			totalRows += n
+			log.Printf("  %-32s %d rows", reflect.TypeOf(m).Elem().Name(), n)
 		}
-		totalRows += n
-		log.Printf("  %-32s %d rows", reflect.TypeOf(m).Elem().Name(), n)
+		return nil
+	})
+	if txErr != nil {
+		return txErr
 	}
 
+	// setval is never rolled back by PostgreSQL, so sequences are resynced only
+	// after the transaction has committed.
 	if err := resetPostgresSequences(dst); err != nil {
 		log.Printf("warning: failed to reset some postgres sequences: %v", err)
 	}
@@ -301,3 +314,32 @@ func tableWithIdColumn(db *gorm.DB, m any) (string, bool) {
 	}
 	return stmt.Table, true
 }
+
+// PrepareSQLiteForMigration rejects SQLite files that are not a panel database
+// before the caller causes any downtime, then AutoMigrates the panel schema
+// onto the file so backups from older versions gain the newer tables and
+// columns the row copy reads. Data-level upgrades are not needed here: they
+// run dialect-agnostically on the destination via InitDB after the import.
+func PrepareSQLiteForMigration(dbPath string) error {
+	gdb, err := gorm.Open(sqlite.Open(dbPath+"?_busy_timeout=10000"), &gorm.Config{Logger: logger.Discard})
+	if err != nil {
+		return err
+	}
+	sqlDB, err := gdb.DB()
+	if err != nil {
+		return err
+	}
+	defer sqlDB.Close()
+
+	for _, table := range []string{"users", "settings", "inbounds"} {
+		if !sqliteTableExists(sqlDB, table) {
+			return fmt.Errorf("not a 3x-ui panel database: required table %q is missing", table)
+		}
+	}
+	for _, m := range migrationModels() {
+		if err := gdb.AutoMigrate(m); err != nil && !isIgnorableDuplicateColumnErr(gdb, err, m) {
+			return fmt.Errorf("upgrade panel schema for %T: %w", m, err)
+		}
+	}
+	return nil
+}

+ 70 - 0
internal/database/migrate_data_test.go

@@ -137,3 +137,73 @@ func TestMigrateData_PreservesFalseDefaultedColumns(t *testing.T) {
 		t.Fatalf("disabled node re-enabled after migration")
 	}
 }
+
+func TestMigrateData_FailedCopyLeavesDestinationUntouched(t *testing.T) {
+	dsn := os.Getenv("XUI_TEST_PG_DSN")
+	if dsn == "" {
+		t.Skip("set XUI_TEST_PG_DSN to a reachable Postgres to run this test")
+	}
+
+	seedSource := func(username string) string {
+		t.Helper()
+		srcPath := t.TempDir() + "/x-ui.db"
+		src, err := gorm.Open(sqlite.Open(srcPath), &gorm.Config{Logger: logger.Discard})
+		if err != nil {
+			t.Fatalf("open sqlite: %v", err)
+		}
+		for _, m := range migrationModels() {
+			if err := src.AutoMigrate(m); err != nil {
+				t.Fatalf("automigrate %T: %v", m, err)
+			}
+		}
+		if err := src.Create(&model.User{Username: username, Password: "pw"}).Error; err != nil {
+			t.Fatalf("seed user: %v", err)
+		}
+		if sqlDB, err := src.DB(); err == nil {
+			sqlDB.Close()
+		}
+		return srcPath
+	}
+
+	dst, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard})
+	if err != nil {
+		t.Fatalf("open postgres: %v", err)
+	}
+	if err := dst.Migrator().DropTable(migrationModels()...); err != nil {
+		t.Fatalf("drop tables: %v", err)
+	}
+
+	if err := MigrateData(seedSource("keep-me"), dsn); err != nil {
+		t.Fatalf("seed destination via MigrateData: %v", err)
+	}
+
+	brokenSrc := seedSource("evil")
+	breaker, err := gorm.Open(sqlite.Open(brokenSrc), &gorm.Config{Logger: logger.Discard})
+	if err != nil {
+		t.Fatalf("reopen broken source: %v", err)
+	}
+	if err := breaker.Exec("DROP TABLE outbound_subscriptions").Error; err != nil {
+		t.Fatalf("drop table from broken source: %v", err)
+	}
+	if sqlDB, err := breaker.DB(); err == nil {
+		sqlDB.Close()
+	}
+
+	if err := MigrateData(brokenSrc, dsn); err == nil {
+		t.Fatal("MigrateData succeeded on a source missing outbound_subscriptions, want error")
+	}
+
+	var keepMe, evil int64
+	if err := dst.Model(&model.User{}).Where("username = ?", "keep-me").Count(&keepMe).Error; err != nil {
+		t.Fatalf("count keep-me: %v", err)
+	}
+	if err := dst.Model(&model.User{}).Where("username = ?", "evil").Count(&evil).Error; err != nil {
+		t.Fatalf("count evil: %v", err)
+	}
+	if keepMe != 1 {
+		t.Fatalf("previous destination data lost after failed migration: keep-me count = %d, want 1", keepMe)
+	}
+	if evil != 0 {
+		t.Fatalf("failed migration leaked partial rows: evil count = %d, want 0", evil)
+	}
+}

+ 29 - 0
internal/database/migration_models_test.go

@@ -0,0 +1,29 @@
+package database
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestMigrationModelsMatchPanelModels(t *testing.T) {
+	names := func(models []any) map[string]bool {
+		set := make(map[string]bool, len(models))
+		for _, m := range models {
+			set[reflect.TypeOf(m).Elem().Name()] = true
+		}
+		return set
+	}
+	panel := names(allModels())
+	migration := names(migrationModels())
+
+	for name := range panel {
+		if !migration[name] {
+			t.Errorf("model %s is in allModels but missing from migrationModels: cross-db migration silently drops its rows", name)
+		}
+	}
+	for name := range migration {
+		if !panel[name] {
+			t.Errorf("model %s is in migrationModels but missing from allModels: its table never exists on a live panel", name)
+		}
+	}
+}

+ 66 - 0
internal/database/prepare_sqlite_test.go

@@ -0,0 +1,66 @@
+package database
+
+import (
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+func TestPrepareSQLiteForMigration(t *testing.T) {
+	t.Run("rejects non-panel database", func(t *testing.T) {
+		dbPath := filepath.Join(t.TempDir(), "random.db")
+		gdb, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{Logger: logger.Discard})
+		if err != nil {
+			t.Fatalf("open sqlite: %v", err)
+		}
+		if err := gdb.Exec("CREATE TABLE notes(id integer primary key, body text)").Error; err != nil {
+			t.Fatalf("create table: %v", err)
+		}
+		closeGorm(gdb)
+
+		err = PrepareSQLiteForMigration(dbPath)
+		if err == nil {
+			t.Fatal("PrepareSQLiteForMigration accepted a non-panel database, want error")
+		}
+		if !strings.Contains(err.Error(), "not a 3x-ui panel database") {
+			t.Fatalf("error = %q, want it to contain %q", err.Error(), "not a 3x-ui panel database")
+		}
+	})
+
+	t.Run("upgrades old panel schema", func(t *testing.T) {
+		dbPath := filepath.Join(t.TempDir(), "old.db")
+		gdb, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{Logger: logger.Discard})
+		if err != nil {
+			t.Fatalf("open sqlite: %v", err)
+		}
+		if err := gdb.AutoMigrate(&model.User{}, &model.Setting{}, &model.Inbound{}); err != nil {
+			t.Fatalf("automigrate legacy subset: %v", err)
+		}
+		closeGorm(gdb)
+
+		if err := PrepareSQLiteForMigration(dbPath); err != nil {
+			t.Fatalf("PrepareSQLiteForMigration rejected an old panel database: %v", err)
+		}
+
+		check, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{Logger: logger.Discard})
+		if err != nil {
+			t.Fatalf("reopen sqlite: %v", err)
+		}
+		defer closeGorm(check)
+		sqlDB, err := check.DB()
+		if err != nil {
+			t.Fatalf("sql db: %v", err)
+		}
+		for _, table := range []string{"client_groups", "client_global_traffics", "outbound_subscriptions"} {
+			if !sqliteTableExists(sqlDB, table) {
+				t.Errorf("table %s was not created by the schema upgrade", table)
+			}
+		}
+	})
+}