Better support Count in chain

This commit is contained in:
Jinzhu 2020-07-01 10:19:52 +08:00
parent 9d7df71332
commit d342f4122a
2 changed files with 10 additions and 0 deletions

View File

@ -269,6 +269,7 @@ func (db *DB) Count(count *int64) (tx *DB) {
if len(tx.Statement.Selects) == 0 { if len(tx.Statement.Selects) == 0 {
tx.Statement.AddClause(clause.Select{Expression: clause.Expr{SQL: "count(1)"}}) tx.Statement.AddClause(clause.Select{Expression: clause.Expr{SQL: "count(1)"}})
defer tx.Statement.AddClause(clause.Select{})
} else if !strings.Contains(strings.ToLower(tx.Statement.Selects[0]), "count(") { } else if !strings.Contains(strings.ToLower(tx.Statement.Selects[0]), "count(") {
expr := clause.Expr{SQL: "count(1)"} expr := clause.Expr{SQL: "count(1)"}
@ -281,6 +282,7 @@ func (db *DB) Count(count *int64) (tx *DB) {
} }
tx.Statement.AddClause(clause.Select{Expression: expr}) tx.Statement.AddClause(clause.Select{Expression: expr})
defer tx.Statement.AddClause(clause.Select{})
} }
tx.Statement.Dest = count tx.Statement.Dest = count

View File

@ -27,6 +27,14 @@ func TestCount(t *testing.T) {
t.Errorf("Count() method should get correct value, expect: %v, got %v", count, len(users)) t.Errorf("Count() method should get correct value, expect: %v, got %v", count, len(users))
} }
if err := DB.Model(&User{}).Where("name = ?", user1.Name).Or("name = ?", user3.Name).Count(&count).Find(&users).Error; err != nil {
t.Errorf(fmt.Sprintf("Count should work, but got err %v", err))
}
if count != int64(len(users)) {
t.Errorf("Count() method should get correct value, expect: %v, got %v", count, len(users))
}
DB.Model(&User{}).Where("name = ?", user1.Name).Count(&count1).Or("name in ?", []string{user2.Name, user3.Name}).Count(&count2) DB.Model(&User{}).Where("name = ?", user1.Name).Count(&count1).Or("name in ?", []string{user2.Name, user3.Name}).Count(&count2)
if count1 != 1 || count2 != 3 { if count1 != 1 || count2 != 3 {
t.Errorf("multiple count in chain should works") t.Errorf("multiple count in chain should works")