Merge pull request #17 from soranoba/fix/with-preload

Fix: TotalCount is always zero when using Preload
soranoba · Feb 27, 2022 · 66849ba · 66849ba
2 parents 83095b5 + cdf0a4f
commit 66849ba
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 1 deletion.
diff --git a/pager.go b/pager.go
@@ -85,8 +85,16 @@ func pagerHandleBeforeQuery(db *gorm.DB) {
 		}
 	}
 	tx.Statement.Clauses = newClauses
-	tx.Model(db.Statement.Dest).Count(&pager.totalCount)
 
+	// NOTE: preload must be deleted.
+	preloads := tx.Statement.Preloads
+	tx.Statement.Preloads = map[string][]interface{}{}
+
+	if err := tx.Model(db.Statement.Dest).Count(&pager.totalCount).Error; err != nil {
+		panic(err)
+	}
+
+	tx.Statement.Preloads = preloads
 	tx.Statement.Clauses = clauses
 }
 

diff --git a/tests/cursor_test.go b/tests/cursor_test.go
@@ -1334,6 +1334,37 @@ func TestCursor_where_clause_is_ambiguous(t *testing.T) {
 	})
 }
 
+func TestCursor_preload(t *testing.T) {
+	db := openDB()
+
+	assertNoError(t, db.Migrator().DropTable(&User{}))
+	assertNoError(t, db.Migrator().DropTable(&Group{}))
+	assertNoError(t, db.AutoMigrate(&Group{}))
+	assertNoError(t, db.AutoMigrate(&User{}))
+
+	assertNoError(
+		t,
+		db.Create(&User{
+			Name: "Alice",
+			Groups: []Group{
+				{Name: "A"},
+				{Name: "B"},
+			},
+		}).Error,
+	)
+
+	var users []*User
+	cursor := &pageboy.Cursor{
+		Limit: 2,
+	}
+
+	assertNoError(t, db.Preload("Groups").Scopes(cursor.Paginate("ID").Order(ASC).Scope()).Find(&users).Error)
+	assertEqual(t, len(users), 1)
+	assertEqual(t, len(users[0].Groups), 2)
+	assertEqual(t, cursor.GetNextAfter(), pbc.FormatCursorString(&users[0].ID))
+	assertEqual(t, cursor.GetNextBefore(), pbc.FormatCursorString(&users[0].ID))
+}
+
 func TestCursor_concurrency(t *testing.T) {
 	db := openDB()
 	assertNoError(t, db.Migrator().DropTable(&cursorModel{}))

diff --git a/tests/pageboy_test.go b/tests/pageboy_test.go
@@ -17,6 +17,18 @@ import (
 	"gorm.io/gorm"
 )
 
+type Group struct {
+	ID     uint `gorm:"primarykey"`
+	Name   string
+	UserID uint
+}
+
+type User struct {
+	ID     uint `gorm:"primarykey"`
+	Name   string
+	Groups []Group
+}
+
 func openDB() *gorm.DB {
 	var (
 		db  *gorm.DB

diff --git a/tests/pager_test.go b/tests/pager_test.go
@@ -130,6 +130,33 @@ func TestPagerPaginateWithWhere(t *testing.T) {
 	assertEqual(t, *pager.Summary(), pageboy.PagerSummary{Page: 1, PerPage: 2, TotalCount: 2, TotalPage: 1})
 }
 
+func TestPager_preload(t *testing.T) {
+	db := openDB()
+
+	assertNoError(t, db.Migrator().DropTable(&User{}))
+	assertNoError(t, db.Migrator().DropTable(&Group{}))
+	assertNoError(t, db.AutoMigrate(&Group{}))
+	assertNoError(t, db.AutoMigrate(&User{}))
+
+	assertNoError(
+		t,
+		db.Create(&User{
+			Name: "Alice",
+			Groups: []Group{
+				{Name: "A"},
+				{Name: "B"},
+			},
+		}).Error,
+	)
+
+	var users []*User
+	pager := &pageboy.Pager{Page: 1, PerPage: 2}
+	assertNoError(t, db.Preload("Groups").Scopes(pager.Scope()).Order("id ASC").Find(&users).Error)
+	assertEqual(t, len(users), 1)
+	assertEqual(t, len(users[0].Groups), 2)
+	assertEqual(t, *pager.Summary(), pageboy.PagerSummary{Page: 1, PerPage: 2, TotalCount: 1, TotalPage: 1})
+}
+
 func TestPager_concurrency(t *testing.T) {
 	db := openDB()
 	assertNoError(t, db.Migrator().DropTable(&cursorModel{}))