Skip to content

Commit c8ad426

Browse files
committed
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.
1 parent 4f597a0 commit c8ad426

2 files changed

Lines changed: 94 additions & 12 deletions

File tree

database/migrate_data.go

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"os"
88
"path"
99
"reflect"
10+
"strings"
1011
"time"
1112

1213
"github.com/mhsanaei/3x-ui/v3/database/model"
@@ -103,25 +104,42 @@ func MigrateData(srcPath, dstDSN string) error {
103104
return nil
104105
}
105106

106-
// copyTable streams every row of `mdl` from src to dst in batches.
107107
func copyTable(src, dst *gorm.DB, mdl any) (int, error) {
108+
const batchSize = 500
109+
108110
sliceType := reflect.SliceOf(reflect.PointerTo(reflect.TypeOf(mdl).Elem()))
109-
batchPtr := reflect.New(sliceType)
110-
batchPtr.Elem().Set(reflect.MakeSlice(sliceType, 0, 0))
111+
112+
// Resolve primary-key columns so paging is deterministic across successive
113+
// LIMIT/OFFSET reads. The model set is trusted (not user input).
114+
stmt := &gorm.Statement{DB: src}
115+
if err := stmt.Parse(mdl); err != nil {
116+
return 0, err
117+
}
118+
order := strings.Join(stmt.Schema.PrimaryFieldDBNames, ", ")
111119

112120
total := 0
113-
err := src.Model(mdl).FindInBatches(batchPtr.Interface(), 500, func(tx *gorm.DB, _ int) error {
114-
batch := batchPtr.Elem()
115-
if batch.Len() == 0 {
116-
return nil
121+
for offset := 0; ; offset += batchSize {
122+
batchPtr := reflect.New(sliceType)
123+
q := src.Model(mdl).Limit(batchSize).Offset(offset)
124+
if order != "" {
125+
q = q.Order(order)
126+
}
127+
if err := q.Find(batchPtr.Interface()).Error; err != nil {
128+
return total, err
129+
}
130+
n := batchPtr.Elem().Len()
131+
if n == 0 {
132+
break
117133
}
118134
if err := dst.CreateInBatches(batchPtr.Interface(), 200).Error; err != nil {
119-
return err
135+
return total, err
136+
}
137+
total += n
138+
if n < batchSize {
139+
break
120140
}
121-
total += batch.Len()
122-
return nil
123-
}).Error
124-
return total, err
141+
}
142+
return total, nil
125143
}
126144

127145
// resetPostgresSequences advances each migrated table's id sequence past MAX(id),

database/migrate_data_test.go

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package database
2+
3+
import (
4+
"os"
5+
"testing"
6+
7+
"github.com/mhsanaei/3x-ui/v3/database/model"
8+
9+
"gorm.io/driver/postgres"
10+
"gorm.io/driver/sqlite"
11+
"gorm.io/gorm"
12+
"gorm.io/gorm/logger"
13+
)
14+
15+
func TestMigrateData_CompositeKeyTableLargerThanBatch(t *testing.T) {
16+
dsn := os.Getenv("XUI_TEST_PG_DSN")
17+
if dsn == "" {
18+
t.Skip("set XUI_TEST_PG_DSN to a reachable Postgres to run this test")
19+
}
20+
21+
// Seed a SQLite source with the full schema and >500 client_inbounds rows.
22+
srcPath := t.TempDir() + "/x-ui.db"
23+
src, err := gorm.Open(sqlite.Open(srcPath), &gorm.Config{Logger: logger.Discard})
24+
if err != nil {
25+
t.Fatalf("open sqlite: %v", err)
26+
}
27+
for _, m := range migrationModels() {
28+
if err := src.AutoMigrate(m); err != nil {
29+
t.Fatalf("automigrate %T: %v", m, err)
30+
}
31+
}
32+
const n = 600 // > batchSize (500) so the between-batches path is exercised
33+
links := make([]model.ClientInbound, 0, n)
34+
for i := 1; i <= n; i++ {
35+
links = append(links, model.ClientInbound{ClientId: i, InboundId: 1})
36+
}
37+
if err := src.CreateInBatches(links, 200).Error; err != nil {
38+
t.Fatalf("seed client_inbounds: %v", err)
39+
}
40+
if sqlDB, err := src.DB(); err == nil {
41+
sqlDB.Close() // flush before MigrateData reopens the file
42+
}
43+
44+
// Make the test re-runnable: drop any tables from a previous run.
45+
dst, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Discard})
46+
if err != nil {
47+
t.Fatalf("open postgres: %v", err)
48+
}
49+
if err := dst.Migrator().DropTable(migrationModels()...); err != nil {
50+
t.Fatalf("drop tables: %v", err)
51+
}
52+
53+
if err := MigrateData(srcPath, dsn); err != nil {
54+
t.Fatalf("MigrateData: %v", err) // fails here before the fix
55+
}
56+
57+
var got int64
58+
if err := dst.Model(&model.ClientInbound{}).Count(&got).Error; err != nil {
59+
t.Fatalf("count: %v", err)
60+
}
61+
if got != n {
62+
t.Fatalf("client_inbounds rows = %d, want %d", got, n)
63+
}
64+
}

0 commit comments

Comments
 (0)