Fix PostgreSQL migration resume and home crash guard

internal/database/database.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"sort"
 	"strings"
 	"time"
 
@@ -18,6 +19,7 @@ import (
 	"gorm.io/gorm"
 	"gorm.io/gorm/clause"
 	"gorm.io/gorm/logger"
+	"gorm.io/gorm/schema"
 
 	"github.com/ShukeBta/MediaStationGo/internal/config"
 	"github.com/ShukeBta/MediaStationGo/internal/model"
@@ -104,11 +106,11 @@ func databaseDialector(cfg *config.Config, dialect string) (gorm.Dialector, erro
 	}
 }
 
-// MigrateSQLiteToCurrentIfNeeded copies an existing SQLite database into a new
-// PostgreSQL database once. It is intentionally conservative: it only runs when
-// the current DB is PostgreSQL, the configured SQLite file exists, and the
-// target business tables are empty. Redis is not migrated because it is a
-// rebuildable cache, not a source of truth.
+// MigrateSQLiteToCurrentIfNeeded copies an existing SQLite database into
+// PostgreSQL. Redis is not migrated because it is a rebuildable cache, not a
+// source of truth.
+const sqliteMigrationCompleteSettingKey = "database.sqlite_migration_complete"
+
 func MigrateSQLiteToCurrentIfNeeded(cfg *config.Config, target *gorm.DB, log *zap.Logger) error {
 	if cfg == nil || target == nil || target.Dialector == nil || target.Dialector.Name() != "postgres" {
 		return nil
@@ -123,11 +125,11 @@ func MigrateSQLiteToCurrentIfNeeded(cfg *config.Config, target *gorm.DB, log *za
 		}
 		return fmt.Errorf("stat sqlite migration source: %w", err)
 	}
-	if empty, err := targetBusinessTablesEmpty(target); err != nil {
+	if complete, err := sqliteMigrationMarkedComplete(target); err != nil {
 		return err
-	} else if !empty {
+	} else if complete {
 		if log != nil {
-			log.Info("skip sqlite to postgres migration: target database already has data")
+			log.Info("skip sqlite to postgres migration: already completed")
 		}
 		return nil
 	}
@@ -150,7 +152,10 @@ func MigrateSQLiteToCurrentIfNeeded(cfg *config.Config, target *gorm.DB, log *za
 	if err != nil {
 		return err
 	}
-	if copied > 0 && log != nil {
+	if err := markSQLiteMigrationComplete(target); err != nil {
+		return err
+	}
+	if log != nil {
 		log.Info("sqlite data migrated to postgres",
 			zap.String("source", sqlitePath),
 			zap.Int64("rows", copied),
@@ -159,21 +164,28 @@ func MigrateSQLiteToCurrentIfNeeded(cfg *config.Config, target *gorm.DB, log *za
 	return nil
 }
 
-func targetBusinessTablesEmpty(db *gorm.DB) (bool, error) {
-	for _, m := range model.AllModels() {
-		table, err := modelTableName(db, m)
-		if err != nil {
-			return false, err
-		}
-		var count int64
-		if err := db.Raw("SELECT COUNT(1) FROM " + quoteIdent(table)).Scan(&count).Error; err != nil {
-			return false, fmt.Errorf("count target table %s: %w", table, err)
-		}
-		if count > 0 {
-			return false, nil
-		}
+func sqliteMigrationMarkedComplete(db *gorm.DB) (bool, error) {
+	var value string
+	err := db.Raw("SELECT value FROM "+quoteIdent("settings")+" WHERE "+quoteIdent("key")+" = ?", sqliteMigrationCompleteSettingKey).Scan(&value).Error
+	if err != nil {
+		return false, fmt.Errorf("check sqlite migration marker: %w", err)
 	}
-	return true, nil
+	return strings.EqualFold(strings.TrimSpace(value), "true"), nil
+}
+
+func markSQLiteMigrationComplete(db *gorm.DB) error {
+	now := time.Now()
+	if err := db.Clauses(clause.OnConflict{
+		Columns:   []clause.Column{{Name: "key"}},
+		DoUpdates: clause.AssignmentColumns([]string{"value", "updated_at"}),
+	}).Create(&model.Setting{
+		Key:       sqliteMigrationCompleteSettingKey,
+		Value:     "true",
+		UpdatedAt: now,
+	}).Error; err != nil {
+		return fmt.Errorf("mark sqlite migration complete: %w", err)
+	}
+	return nil
 }
 
 func copyModelTables(src, target *gorm.DB, batchSize int) (int64, error) {
@@ -186,6 +198,10 @@ func copyModelTables(src, target *gorm.DB, batchSize int) (int64, error) {
 		if err != nil {
 			return copied, err
 		}
+		primaryColumns, err := modelPrimaryColumns(src, m)
+		if err != nil {
+			return copied, fmt.Errorf("inspect model %T primary keys: %w", m, err)
+		}
 		exists, err := sqliteTableExists(src, table)
 		if err != nil {
 			return copied, err
@@ -204,10 +220,6 @@ func copyModelTables(src, target *gorm.DB, batchSize int) (int64, error) {
 		if err := target.Raw("SELECT COUNT(1) FROM " + quoteIdent(table)).Scan(&targetCount).Error; err != nil {
 			return copied, fmt.Errorf("count target table %s: %w", table, err)
 		}
-		if targetCount > 0 {
-			continue
-		}
-
 		modelType := reflect.TypeOf(m)
 		if modelType.Kind() != reflect.Ptr {
 			return copied, fmt.Errorf("model %T is not a pointer", m)
@@ -217,17 +229,156 @@ func copyModelTables(src, target *gorm.DB, batchSize int) (int64, error) {
 		if err := src.Unscoped().Find(slicePtr.Interface()).Error; err != nil {
 			return copied, fmt.Errorf("read sqlite table %s: %w", table, err)
 		}
-		if slicePtr.Elem().Len() == 0 {
+		filtered := slicePtr.Elem()
+		if targetCount > 0 {
+			primaryKeySet, err := targetPrimaryKeySet(target, table, primaryColumns)
+			if err != nil {
+				return copied, err
+			}
+			filtered = filterRowsMissingInTarget(target, table, primaryColumns, filtered, primaryKeySet)
+		}
+		if filtered.Len() == 0 {
 			continue
 		}
-		if err := target.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(slicePtr.Interface(), batchSize).Error; err != nil {
+		filteredPtr := reflect.New(filtered.Type())
+		filteredPtr.Elem().Set(filtered)
+		if err := target.Clauses(clause.OnConflict{DoNothing: true}).CreateInBatches(filteredPtr.Interface(), batchSize).Error; err != nil {
 			return copied, fmt.Errorf("copy sqlite table %s: %w", table, err)
 		}
-		copied += int64(slicePtr.Elem().Len())
+		copied += int64(filtered.Len())
 	}
 	return copied, nil
 }
 
+func modelPrimaryColumns(db *gorm.DB, m any) ([]string, error) {
+	stmt := &gorm.Statement{DB: db}
+	if err := stmt.Parse(m); err != nil {
+		return nil, err
+	}
+	var cols []string
+	for _, field := range stmt.Schema.PrimaryFields {
+		cols = append(cols, field.DBName)
+	}
+	if len(cols) == 0 {
+		return nil, fmt.Errorf("no primary key columns")
+	}
+	return cols, nil
+}
+
+func targetPrimaryKeySet(target *gorm.DB, table string, primaryColumns []string) (map[string]struct{}, error) {
+	if len(primaryColumns) != 1 {
+		return nil, nil
+	}
+	var values []string
+	if err := target.Raw("SELECT " + quoteIdent(primaryColumns[0]) + " FROM " + quoteIdent(table)).Scan(&values).Error; err != nil {
+		return nil, fmt.Errorf("read target primary keys for table %s: %w", table, err)
+	}
+	set := make(map[string]struct{}, len(values))
+	for _, value := range values {
+		set[value] = struct{}{}
+	}
+	return set, nil
+}
+
+func filterRowsMissingInTarget(target *gorm.DB, table string, primaryColumns []string, rows reflect.Value, primaryKeySet map[string]struct{}) reflect.Value {
+	if rows.Kind() != reflect.Slice || rows.Len() == 0 || len(primaryColumns) == 0 {
+		return rows
+	}
+	out := reflect.MakeSlice(rows.Type(), 0, rows.Len())
+	for i := 0; i < rows.Len(); i++ {
+		row := rows.Index(i)
+		keys, ok := rowPrimaryKeys(row, primaryColumns)
+		if !ok {
+			out = reflect.Append(out, row)
+			continue
+		}
+		if primaryKeySet != nil {
+			if _, exists := primaryKeySet[fmt.Sprint(keys[primaryColumns[0]])]; !exists {
+				out = reflect.Append(out, row)
+			}
+			continue
+		}
+		if !targetHasPrimaryKey(target, table, keys) {
+			out = reflect.Append(out, row)
+		}
+	}
+	return out
+}
+
+func rowPrimaryKeys(row reflect.Value, primaryColumns []string) (map[string]any, bool) {
+	if row.Kind() == reflect.Pointer {
+		if row.IsNil() {
+			return nil, false
+		}
+		row = row.Elem()
+	}
+	if row.Kind() != reflect.Struct {
+		return nil, false
+	}
+	keys := make(map[string]any, len(primaryColumns))
+	for _, column := range primaryColumns {
+		value, ok := fieldByDBName(row, column)
+		if !ok || value.IsZero() {
+			return nil, false
+		}
+		keys[column] = value.Interface()
+	}
+	return keys, true
+}
+
+func fieldByDBName(row reflect.Value, column string) (reflect.Value, bool) {
+	rowType := row.Type()
+	for i := 0; i < row.NumField(); i++ {
+		fieldType := rowType.Field(i)
+		field := row.Field(i)
+		if fieldType.Anonymous {
+			if value, ok := fieldByDBName(field, column); ok {
+				return value, true
+			}
+		}
+		if columnNameForStructField(fieldType) == column {
+			if field.Kind() == reflect.Pointer && field.IsNil() {
+				return reflect.Value{}, false
+			}
+			return field, field.CanInterface()
+		}
+	}
+	return reflect.Value{}, false
+}
+
+func columnNameForStructField(field reflect.StructField) string {
+	if field.PkgPath != "" && !field.Anonymous {
+		return ""
+	}
+	tag := field.Tag.Get("gorm")
+	settings := schema.ParseTagSetting(tag, ";")
+	if column := settings["COLUMN"]; column != "" {
+		return column
+	}
+	return schema.NamingStrategy{}.ColumnName("", field.Name)
+}
+
+func targetHasPrimaryKey(target *gorm.DB, table string, keys map[string]any) bool {
+	where := make([]string, 0, len(keys))
+	args := make([]any, 0, len(keys))
+	for _, column := range sortedMapKeys(keys) {
+		where = append(where, quoteIdent(column)+" = ?")
+		args = append(args, keys[column])
+	}
+	var count int64
+	err := target.Raw("SELECT COUNT(1) FROM "+quoteIdent(table)+" WHERE "+strings.Join(where, " AND "), args...).Scan(&count).Error
+	return err == nil && count > 0
+}
+
+func sortedMapKeys(m map[string]any) []string {
+	keys := make([]string, 0, len(m))
+	for key := range m {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	return keys
+}
+
 func sqliteTableExists(db *gorm.DB, table string) (bool, error) {
 	var count int64
 	if err := db.Raw(`SELECT COUNT(1) FROM sqlite_master WHERE type = 'table' AND name = ?`, table).Scan(&count).Error; err != nil {
@@ -341,6 +492,9 @@ func AutoMigrate(db *gorm.DB) error {
 	if err := db.AutoMigrate(model.AllModels()...); err != nil {
 		return err
 	}
+	if err := ensurePostgresColumnCompatibility(db); err != nil {
+		return err
+	}
 	if err := enforceTelegramBindingOneToOne(db); err != nil {
 		return err
 	}
@@ -353,6 +507,21 @@ func AutoMigrate(db *gorm.DB) error {
 	return nil
 }
 
+func ensurePostgresColumnCompatibility(db *gorm.DB) error {
+	if !isPostgres(db) {
+		return nil
+	}
+	statements := []string{
+		`ALTER TABLE media ALTER COLUMN container TYPE varchar(128)`,
+	}
+	for _, stmt := range statements {
+		if err := db.Exec(stmt).Error; err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func ensurePerformanceIndexes(db *gorm.DB) error {
 	statements := []string{
 		`CREATE INDEX IF NOT EXISTS idx_media_library_created_active ON media(library_id, created_at DESC) WHERE deleted_at IS NULL`,
@@ -386,6 +555,10 @@ func isSQLite(db *gorm.DB) bool {
 	return db != nil && db.Dialector != nil && db.Dialector.Name() == "sqlite"
 }
 
+func isPostgres(db *gorm.DB) bool {
+	return db != nil && db.Dialector != nil && db.Dialector.Name() == "postgres"
+}
+
 // mediaSearchIndexSchemaVersion 标识 FTS 索引的物理布局版本。
 // v2：FTS 行的 rowid 与 media.rowid 对齐，并由触发器实时维护。
 const mediaSearchIndexSchemaVersion = 2