Merge pull request uptrace#343 from uptrace/fix/scan-and-count-tx

fix: change ScanAndCount to work with transactions

internal/dbtest/db_test.go

@@ -235,6 +235,7 @@ func TestDB(t *testing.T) {
 		{testBinaryData},
 		{testUpsert},
 		{testMultiUpdate},
+		{testTxScanAndCount},
 	}
 
 	testEachDB(t, func(t *testing.T, dbName string, db *bun.DB) {
@@ -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)
+	}
+}

query_select.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
-	"sync"
 
 	"github.com/uptrace/bun/internal"
 	"github.com/uptrace/bun/schema"
@@ -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
 }