Skip to content

Commit

Permalink
fix: change ScanAndCount to work with transactions
Browse files Browse the repository at this point in the history
  • Loading branch information
vmihailenco committed Dec 3, 2021
1 parent 47e28c3 commit 5b3f2c0
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 31 deletions.
22 changes: 22 additions & 0 deletions internal/dbtest/db_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,7 @@ func TestDB(t *testing.T) {
{testBinaryData},
{testUpsert},
{testMultiUpdate},
{testTxScanAndCount},
}

testEachDB(t, func(t *testing.T, dbName string, db *bun.DB) {
Expand Down Expand Up @@ -1007,3 +1008,24 @@ func testMultiUpdate(t *testing.T, db *bun.DB) {
Exec(ctx)
require.NoError(t, err)
}

func testTxScanAndCount(t *testing.T, db *bun.DB) {
type Model struct {
ID int64
Str string
}

ctx := context.Background()

err := db.ResetModel(ctx, (*Model)(nil))
require.NoError(t, err)

for i := 0; i < 100; i++ {
err := db.RunInTx(ctx, nil, func(ctx context.Context, tx bun.Tx) error {
var models []Model
_, err := tx.NewSelect().Model(&models).ScanAndCount(ctx)
return err
})
require.NoError(t, err)
}
}
36 changes: 5 additions & 31 deletions query_select.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ import (
"fmt"
"strconv"
"strings"
"sync"

"github.com/uptrace/bun/internal"
"github.com/uptrace/bun/schema"
Expand Down Expand Up @@ -782,42 +781,17 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) {
}

func (q *SelectQuery) ScanAndCount(ctx context.Context, dest ...interface{}) (int, error) {
var count int
var wg sync.WaitGroup
var mu sync.Mutex
var firstErr error

if q.limit >= 0 {
wg.Add(1)
go func() {
defer wg.Done()

if err := q.Scan(ctx, dest...); err != nil {
mu.Lock()
if firstErr == nil {
firstErr = err
}
mu.Unlock()
}
}()
firstErr = q.Scan(ctx, dest...)
}

wg.Add(1)
go func() {
defer wg.Done()

var err error
count, err = q.Count(ctx)
if err != nil {
mu.Lock()
if firstErr == nil {
firstErr = err
}
mu.Unlock()
}
}()
count, err := q.Count(ctx)
if err != nil && firstErr == nil {
firstErr = err
}

wg.Wait()
return count, firstErr
}

Expand Down

0 comments on commit 5b3f2c0

Please sign in to comment.