A better "dequeue" abstraction defined on reviews table

Author: VojtechVitek

tests/database_test.go

@@ -41,16 +41,6 @@ func (db *Database) WithTx(tx pgx.Tx) *Database {
 	return initDB(pgkitDB)
 }
 
-func (db *Database) Close() { db.DB.Conn.Close() }
-
-type accountsTable struct {
-	*pgkit.Table[Account, *Account, int64]
-}
-
-type articlesTable struct {
-	*pgkit.Table[Article, *Article, uint64]
-}
-
-type reviewsTable struct {
-	*pgkit.Table[Review, *Review, uint64]
+func (db *Database) Close() {
+	db.DB.Conn.Close()
 }

tests/table_test.go

@@ -5,7 +5,6 @@ import (
 	"slices"
 	"sync"
 	"testing"
-	"time"
 
 	sq "github.com/Masterminds/squirrel"
 	"github.com/stretchr/testify/require"
@@ -183,49 +182,28 @@ func TestLockForUpdates(t *testing.T) {
 		err = db.Reviews.Save(ctx, reviews...)
 		require.NoError(t, err, "create review")
 
-		where := sq.Eq{
-			"status":     ReviewStatusPending,
-			"deleted_at": nil,
-		}
-		orderBy := []string{
-			"created_at ASC",
-		}
-		limit := uint64(10)
-
-		var processedIDs [][]uint64 = make([][]uint64, 10)
+		var ids [][]uint64 = make([][]uint64, 10)
 		var wg sync.WaitGroup
 
 		for range 10 {
 			wg.Add(1)
 			go func() {
 				defer wg.Done()
 
-				var processReviews []*Review
+				reviews, err := db.Reviews.DequeueForProcessing(ctx, 10)
+				require.NoError(t, err, "dequeue reviews")
 
-				err := db.Reviews.LockForUpdates(ctx, where, orderBy, limit, func(reviews []*Review) {
-					now := time.Now().UTC()
-					for _, review := range reviews {
-						review.Status = ReviewStatusProcessing
-						review.ProcessedAt = &now
-					}
-
-					processReviews = reviews
-				})
-				require.NoError(t, err, "lock for update")
-
-				for _, review := range processReviews {
+				for i, review := range reviews {
 					go worker.ProcessReview(ctx, review)
-				}
 
-				for i, review := range processReviews {
-					processedIDs[i] = append(processedIDs[i], review.ID)
+					ids[i] = append(ids[i], review.ID)
 				}
 			}()
 		}
 		wg.Wait()
 
 		// Ensure that all reviews were picked up for processing exactly once.
-		uniqueIDs := slices.Concat(processedIDs...)
+		uniqueIDs := slices.Concat(ids...)
 		slices.Sort(uniqueIDs)
 		uniqueIDs = slices.Compact(uniqueIDs)
 		require.Equal(t, 100, len(uniqueIDs), "number of unique reviews picked up for processing should be 100")

tests/tables_test.go

@@ -0,0 +1,47 @@
+package pgkit_test
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	sq "github.com/Masterminds/squirrel"
+	"github.com/goware/pgkit/v2"
+)
+
+type accountsTable struct {
+	*pgkit.Table[Account, *Account, int64]
+}
+
+type articlesTable struct {
+	*pgkit.Table[Article, *Article, uint64]
+}
+
+type reviewsTable struct {
+	*pgkit.Table[Review, *Review, uint64]
+}
+
+func (w *reviewsTable) DequeueForProcessing(ctx context.Context, limit uint64) ([]*Review, error) {
+	var dequeued []*Review
+	where := sq.Eq{
+		"status":     ReviewStatusPending,
+		"deleted_at": nil,
+	}
+	orderBy := []string{
+		"created_at ASC",
+	}
+
+	err := w.LockForUpdates(ctx, where, orderBy, limit, func(reviews []*Review) {
+		now := time.Now().UTC()
+		for _, review := range reviews {
+			review.Status = ReviewStatusProcessing
+			review.ProcessedAt = &now
+		}
+		dequeued = reviews
+	})
+	if err != nil {
+		return nil, fmt.Errorf("lock for updates: %w", err)
+	}
+
+	return dequeued, nil
+}