From c8ad42631c6f65c6b8646d20899968cfd6a0029a Mon Sep 17 00:00:00 2001
From: MHSanaei <ho3ein.sanaei@gmail.com>
Date: Tue, 2 Jun 2026 04:20:42 +0200
Subject: [PATCH] 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.
---
 database/migrate_data.go      | 42 ++++++++++++++++-------
 database/migrate_data_test.go | 64 +++++++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 12 deletions(-)
 create mode 100644 database/migrate_data_test.go

diff --git a/database/migrate_data.go b/database/migrate_data.go
index 0a29ffb4..f6f6e43f 100644
--- a/database/migrate_data.go
+++ b/database/migrate_data.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"path"
 	"reflect"
+	"strings"
 	"time"
 
 	"github.com/mhsanaei/3x-ui/v3/database/model"
@@ -103,25 +104,42 @@ func MigrateData(srcPath, dstDSN string) error {
 	return nil
 }
 
-// copyTable streams every row of `mdl` from src to dst in batches.
 func copyTable(src, dst *gorm.DB, mdl any) (int, error) {
+	const batchSize = 500
+
 	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
-	err := src.Model(mdl).FindInBatches(batchPtr.Interface(), 500, func(tx *gorm.DB, _ int) error {
-		batch := batchPtr.Elem()
-		if batch.Len() == 0 {
-			return nil
+	for offset := 0; ; offset += batchSize {
+		batchPtr := reflect.New(sliceType)
+		q := src.Model(mdl).Limit(batchSize).Offset(offset)
+		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 {
-			return err
+			return total, err
 		}
-		total += batch.Len()
-		return nil
-	}).Error
-	return total, err
+		total += n
+		if n < batchSize {
+			break
+		}
+	}
+	return total, nil
 }
 
 // resetPostgresSequences advances each migrated table's id sequence past MAX(id),
diff --git a/database/migrate_data_test.go b/database/migrate_data_test.go
new file mode 100644
index 00000000..081d77d9
--- /dev/null
+++ b/database/migrate_data_test.go
@@ -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)
+	}
+}