MeBox/internal/database/sqlite_migration_copy.go

package database

import (
	"fmt"
	"reflect"
	"sort"
	"strings"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
	"gorm.io/gorm/schema"

	"github.com/truewhile/MeBox/internal/model"
)

func copyModelTables(src, target *gorm.DB, batchSize int) (map[string]int64, int64, error) {
	if batchSize <= 0 {
		batchSize = 500
	}
	tableCounts := make(map[string]int64)
	var totalCopied int64
	for _, m := range model.AllModels() {
		table, err := modelTableName(src, m)
		if err != nil {
			return tableCounts, totalCopied, err
		}
		primaryColumns, err := modelPrimaryColumns(src, m)
		if err != nil {
			return tableCounts, totalCopied, fmt.Errorf("inspect model %T primary keys: %w", m, err)
		}
		exists, err := sqliteTableExists(src, table)
		if err != nil {
			return tableCounts, totalCopied, err
		}
		if !exists {
			continue
		}
		var sourceCount int64
		if err := src.Raw("SELECT COUNT(1) FROM " + quoteIdent(table)).Scan(&sourceCount).Error; err != nil {
			return tableCounts, totalCopied, fmt.Errorf("count sqlite table %s: %w", table, err)
		}
		if sourceCount == 0 {
			continue
		}
		var targetCount int64
		if err := target.Raw("SELECT COUNT(1) FROM " + quoteIdent(table)).Scan(&targetCount).Error; err != nil {
			return tableCounts, totalCopied, fmt.Errorf("count target table %s: %w", table, err)
		}
		modelType := reflect.TypeOf(m)
		if modelType.Kind() != reflect.Ptr {
			return tableCounts, totalCopied, fmt.Errorf("model %T is not a pointer", m)
		}
		sliceType := reflect.SliceOf(modelType.Elem())
		slicePtr := reflect.New(sliceType)
		if err := src.Unscoped().Find(slicePtr.Interface()).Error; err != nil {
			return tableCounts, totalCopied, fmt.Errorf("read sqlite table %s: %w", table, err)
		}
		filtered := slicePtr.Elem()
		if targetCount > 0 {
			primaryKeySet, err := targetPrimaryKeySet(target, table, primaryColumns)
			if err != nil {
				return tableCounts, totalCopied, err
			}
			filtered = filterRowsMissingInTarget(target, table, primaryColumns, filtered, primaryKeySet)
		}
		if filtered.Len() == 0 {
			continue
		}
		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 tableCounts, totalCopied, fmt.Errorf("copy sqlite table %s: %w", table, err)
		}
		copiedForTable := int64(filtered.Len())
		tableCounts[table] = copiedForTable
		totalCopied += copiedForTable
	}
	return tableCounts, totalCopied, 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 {
		return false, fmt.Errorf("inspect sqlite table %s: %w", table, err)
	}
	return count > 0, nil
}

func modelTableName(db *gorm.DB, m any) (string, error) {
	stmt := &gorm.Statement{DB: db}
	if err := stmt.Parse(m); err != nil {
		return "", err
	}
	return stmt.Schema.Table, nil
}

func quoteIdent(value string) string {
	return `"` + strings.ReplaceAll(value, `"`, `""`) + `"`
}