fix(migrate): copy composite-key tables without FindInBatches (#4787)

SQLite to Postgres migration aborted with "copy *model.ClientInbound:
primary key required" on installs whose client_inbounds table exceeds
one read batch (500 rows). gorm's FindInBatches pages between batches
using a single PrioritizedPrimaryField, which composite-key tables
(client_id + inbound_id, no surrogate id) do not have, so it returns
ErrPrimaryKeyRequired once a table holds more than one batch.

Replace FindInBatches in copyTable with explicit LIMIT/OFFSET paging
ordered by the model's primary-key columns. This works for every table
including composite-key ones, keeps memory bounded, and changes no
schema.

Add a Postgres-gated regression test covering a >500-row composite-key
table.
This commit is contained in:
MHSanaei 2026-06-02 04:20:42 +02:00
parent 4f597a08c4
commit c8ad42631c
No known key found for this signature in database
GPG key ID: 7E4060F2FBE5AB7A
2 changed files with 94 additions and 12 deletions

View file

@ -7,6 +7,7 @@ import (
"os" "os"
"path" "path"
"reflect" "reflect"
"strings"
"time" "time"
"github.com/mhsanaei/3x-ui/v3/database/model" "github.com/mhsanaei/3x-ui/v3/database/model"
@ -103,25 +104,42 @@ func MigrateData(srcPath, dstDSN string) error {
return nil return nil
} }
// copyTable streams every row of `mdl` from src to dst in batches.
func copyTable(src, dst *gorm.DB, mdl any) (int, error) { func copyTable(src, dst *gorm.DB, mdl any) (int, error) {
const batchSize = 500
sliceType := reflect.SliceOf(reflect.PointerTo(reflect.TypeOf(mdl).Elem())) sliceType := reflect.SliceOf(reflect.PointerTo(reflect.TypeOf(mdl).Elem()))
batchPtr := reflect.New(sliceType)
batchPtr.Elem().Set(reflect.MakeSlice(sliceType, 0, 0)) // Resolve primary-key columns so paging is deterministic across successive
// LIMIT/OFFSET reads. The model set is trusted (not user input).
stmt := &gorm.Statement{DB: src}
if err := stmt.Parse(mdl); err != nil {
return 0, err
}
order := strings.Join(stmt.Schema.PrimaryFieldDBNames, ", ")
total := 0 total := 0
err := src.Model(mdl).FindInBatches(batchPtr.Interface(), 500, func(tx *gorm.DB, _ int) error { for offset := 0; ; offset += batchSize {
batch := batchPtr.Elem() batchPtr := reflect.New(sliceType)
if batch.Len() == 0 { q := src.Model(mdl).Limit(batchSize).Offset(offset)
return nil if order != "" {
q = q.Order(order)
}
if err := q.Find(batchPtr.Interface()).Error; err != nil {
return total, err
}
n := batchPtr.Elem().Len()
if n == 0 {
break
} }
if err := dst.CreateInBatches(batchPtr.Interface(), 200).Error; err != nil { if err := dst.CreateInBatches(batchPtr.Interface(), 200).Error; err != nil {
return err return total, err
} }
total += batch.Len() total += n
return nil if n < batchSize {
}).Error break
return total, err }
}
return total, nil
} }
// resetPostgresSequences advances each migrated table's id sequence past MAX(id), // resetPostgresSequences advances each migrated table's id sequence past MAX(id),

View file

@ -0,0 +1,64 @@
package database
import (
"os"
"testing"
"github.com/mhsanaei/3x-ui/v3/database/model"
"gorm.io/driver/postgres"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
func TestMigrateData_CompositeKeyTableLargerThanBatch(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")
}
// Seed a SQLite source with the full schema and >500 client_inbounds rows.
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)
}
}
const n = 600 // > batchSize (500) so the between-batches path is exercised
links := make([]model.ClientInbound, 0, n)
for i := 1; i <= n; i++ {
links = append(links, model.ClientInbound{ClientId: i, InboundId: 1})
}
if err := src.CreateInBatches(links, 200).Error; err != nil {
t.Fatalf("seed client_inbounds: %v", err)
}
if sqlDB, err := src.DB(); err == nil {
sqlDB.Close() // flush before MigrateData reopens the file
}
// Make the test re-runnable: drop any tables from a previous run.
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(srcPath, dsn); err != nil {
t.Fatalf("MigrateData: %v", err) // fails here before the fix
}
var got int64
if err := dst.Model(&model.ClientInbound{}).Count(&got).Error; err != nil {
t.Fatalf("count: %v", err)
}
if got != n {
t.Fatalf("client_inbounds rows = %d, want %d", got, n)
}
}