Merge pull request #10121 from ellemouton/graphPerf6

[5] sqldb+graph/db: add and use new paginate & batch helper

graph/db/graph_test.go

@@ -2510,6 +2510,11 @@ func TestFilterKnownChanIDs(t *testing.T) {
 // methods that acquire the cache mutex along with the DB mutex.
 func TestStressTestChannelGraphAPI(t *testing.T) {
 	t.Parallel()
+
+	if testing.Short() {
+		t.Skipf("Skipping test in short mode")
+	}
+
 	ctx := context.Background()
 
 	graph := MakeTestGraph(t)

graph/db/sql_migration_test.go

@@ -1177,6 +1177,10 @@ func assertResultState(t *testing.T, sql *SQLStore, expState dbState) {
 func TestMigrateGraphToSQLRapid(t *testing.T) {
 	t.Parallel()
 
+	if testing.Short() {
+		t.Skipf("skipping test in short mode")
+	}
+
 	dbFixture := NewTestDBFixture(t)
 
 	rapid.Check(t, func(rt *rapid.T) {

graph/db/sql_store.go

@@ -806,49 +806,18 @@ func (s *SQLStore) ForEachSourceNodeChannel(ctx context.Context,
 func (s *SQLStore) ForEachNode(ctx context.Context,
 	cb func(tx NodeRTx) error, reset func()) error {
 
-	var lastID int64
-
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		nodeCB := func(dbID int64, node *models.LightningNode) error {
+		return forEachNodePaginated(
-			err := cb(newSQLGraphNodeTx(
+			ctx, s.cfg.QueryCfg, db,
-				db, s.cfg.ChainHash, dbID, node,
+			ProtocolV1,
-			))
+			func(ctx context.Context, dbNodeID int64,
-			if err != nil {
+				node *models.LightningNode) error {
-				return fmt.Errorf("callback failed for "+
-					"node(id=%d): %w", dbID, err)
-			}
-			lastID = dbID
 
-			return nil
+				return cb(newSQLGraphNodeTx(
-		}
+					db, s.cfg.ChainHash, dbNodeID, node,
-
+				))
-		for {
+			},
-			nodes, err := db.ListNodesPaginated(
+		)
-				ctx, sqlc.ListNodesPaginatedParams{
-					Version: int16(ProtocolV1),
-					ID:      lastID,
-					Limit:   s.cfg.QueryCfg.MaxPageSize,
-				},
-			)
-			if err != nil {
-				return fmt.Errorf("unable to fetch nodes: %w",
-					err)
-			}
-
-			if len(nodes) == 0 {
-				break
-			}
-
-			err = forEachNodeInBatch(
-				ctx, s.cfg.QueryCfg, db, nodes, nodeCB,
-			)
-			if err != nil {
-				return fmt.Errorf("unable to iterate over "+
-					"nodes: %w", err)
-			}
-		}
-
-		return nil
 	}, reset)
 }
 
@@ -1328,115 +1297,8 @@ func (s *SQLStore) ForEachChannel(ctx context.Context,
 	cb func(*models.ChannelEdgeInfo, *models.ChannelEdgePolicy,
 		*models.ChannelEdgePolicy) error, reset func()) error {
 
-	handleChannel := func(db SQLQueries, batchData *batchChannelData,
-		row sqlc.ListChannelsWithPoliciesPaginatedRow) error {
-
-		node1, node2, err := buildNodeVertices(
-			row.Node1Pubkey, row.Node2Pubkey,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build node vertices: %w",
-				err)
-		}
-
-		edge, err := buildEdgeInfoWithBatchData(
-			s.cfg.ChainHash, row.GraphChannel, node1, node2,
-			batchData,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build channel info: %w",
-				err)
-		}
-
-		dbPol1, dbPol2, err := extractChannelPolicies(row)
-		if err != nil {
-			return fmt.Errorf("unable to extract channel "+
-				"policies: %w", err)
-		}
-
-		p1, p2, err := buildChanPoliciesWithBatchData(
-			dbPol1, dbPol2, edge.ChannelID, node1, node2, batchData,
-		)
-		if err != nil {
-			return fmt.Errorf("unable to build channel "+
-				"policies: %w", err)
-		}
-
-		err = cb(edge, p1, p2)
-		if err != nil {
-			return fmt.Errorf("callback failed for channel "+
-				"id=%d: %w", edge.ChannelID, err)
-		}
-
-		return nil
-	}
-
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		lastID := int64(-1)
+		return forEachChannelWithPolicies(ctx, db, s.cfg, cb)
-		for {
-			//nolint:ll
-			rows, err := db.ListChannelsWithPoliciesPaginated(
-				ctx, sqlc.ListChannelsWithPoliciesPaginatedParams{
-					Version: int16(ProtocolV1),
-					ID:      lastID,
-					Limit:   s.cfg.QueryCfg.MaxPageSize,
-				},
-			)
-			if err != nil {
-				return err
-			}
-
-			if len(rows) == 0 {
-				break
-			}
-
-			// Collect the channel & policy IDs that we want to
-			// do a batch collection for.
-			var (
-				channelIDs = make([]int64, len(rows))
-				policyIDs  = make([]int64, 0, len(rows)*2)
-			)
-			for i, row := range rows {
-				channelIDs[i] = row.GraphChannel.ID
-
-				// Extract policy IDs from the row
-				dbPol1, dbPol2, err := extractChannelPolicies(
-					row,
-				)
-				if err != nil {
-					return fmt.Errorf("unable to extract "+
-						"channel policies: %w", err)
-				}
-
-				if dbPol1 != nil {
-					policyIDs = append(policyIDs, dbPol1.ID)
-				}
-
-				if dbPol2 != nil {
-					policyIDs = append(policyIDs, dbPol2.ID)
-				}
-			}
-
-			batchData, err := batchLoadChannelData(
-				ctx, s.cfg.QueryCfg, db, channelIDs,
-				policyIDs,
-			)
-			if err != nil {
-				return fmt.Errorf("unable to batch load "+
-					"channel data: %w", err)
-			}
-
-			for _, row := range rows {
-				err := handleChannel(db, batchData, row)
-				if err != nil {
-					return err
-				}
-
-				lastID = row.GraphChannel.ID
-			}
-		}
-
-		return nil
 	}, reset)
 }
 
@@ -5082,3 +4944,169 @@ func batchLoadChannelPolicyExtrasHelper(ctx context.Context,
 		},
 	)
 }
+
+// forEachNodePaginated executes a paginated query to process each node in the
+// graph. It uses the provided SQLQueries interface to fetch nodes in batches
+// and applies the provided processNode function to each node.
+func forEachNodePaginated(ctx context.Context, cfg *sqldb.QueryConfig,
+	db SQLQueries, protocol ProtocolVersion,
+	processNode func(context.Context, int64,
+		*models.LightningNode) error) error {
+
+	pageQueryFunc := func(ctx context.Context, lastID int64,
+		limit int32) ([]sqlc.GraphNode, error) {
+
+		return db.ListNodesPaginated(
+			ctx, sqlc.ListNodesPaginatedParams{
+				Version: int16(protocol),
+				ID:      lastID,
+				Limit:   limit,
+			},
+		)
+	}
+
+	extractPageCursor := func(node sqlc.GraphNode) int64 {
+		return node.ID
+	}
+
+	collectFunc := func(node sqlc.GraphNode) (int64, error) {
+		return node.ID, nil
+	}
+
+	batchQueryFunc := func(ctx context.Context,
+		nodeIDs []int64) (*batchNodeData, error) {
+
+		return batchLoadNodeData(ctx, cfg, db, nodeIDs)
+	}
+
+	processItem := func(ctx context.Context, dbNode sqlc.GraphNode,
+		batchData *batchNodeData) error {
+
+		node, err := buildNodeWithBatchData(&dbNode, batchData)
+		if err != nil {
+			return fmt.Errorf("unable to build "+
+				"node(id=%d): %w", dbNode.ID, err)
+		}
+
+		return processNode(ctx, dbNode.ID, node)
+	}
+
+	return sqldb.ExecuteCollectAndBatchWithSharedDataQuery(
+		ctx, cfg, int64(-1), pageQueryFunc, extractPageCursor,
+		collectFunc, batchQueryFunc, processItem,
+	)
+}
+
+// forEachChannelWithPolicies executes a paginated query to process each channel
+// with policies in the graph.
+func forEachChannelWithPolicies(ctx context.Context, db SQLQueries,
+	cfg *SQLStoreConfig, processChannel func(*models.ChannelEdgeInfo,
+		*models.ChannelEdgePolicy,
+		*models.ChannelEdgePolicy) error) error {
+
+	type channelBatchIDs struct {
+		channelID int64
+		policyIDs []int64
+	}
+
+	pageQueryFunc := func(ctx context.Context, lastID int64,
+		limit int32) ([]sqlc.ListChannelsWithPoliciesPaginatedRow,
+		error) {
+
+		return db.ListChannelsWithPoliciesPaginated(
+			ctx, sqlc.ListChannelsWithPoliciesPaginatedParams{
+				Version: int16(ProtocolV1),
+				ID:      lastID,
+				Limit:   limit,
+			},
+		)
+	}
+
+	extractPageCursor := func(
+		row sqlc.ListChannelsWithPoliciesPaginatedRow) int64 {
+
+		return row.GraphChannel.ID
+	}
+
+	collectFunc := func(row sqlc.ListChannelsWithPoliciesPaginatedRow) (
+		channelBatchIDs, error) {
+
+		ids := channelBatchIDs{
+			channelID: row.GraphChannel.ID,
+		}
+
+		// Extract policy IDs from the row.
+		dbPol1, dbPol2, err := extractChannelPolicies(row)
+		if err != nil {
+			return ids, err
+		}
+
+		if dbPol1 != nil {
+			ids.policyIDs = append(ids.policyIDs, dbPol1.ID)
+		}
+		if dbPol2 != nil {
+			ids.policyIDs = append(ids.policyIDs, dbPol2.ID)
+		}
+
+		return ids, nil
+	}
+
+	batchDataFunc := func(ctx context.Context,
+		allIDs []channelBatchIDs) (*batchChannelData, error) {
+
+		// Separate channel IDs from policy IDs.
+		var (
+			channelIDs = make([]int64, len(allIDs))
+			policyIDs  = make([]int64, 0, len(allIDs)*2)
+		)
+
+		for i, ids := range allIDs {
+			channelIDs[i] = ids.channelID
+			policyIDs = append(policyIDs, ids.policyIDs...)
+		}
+
+		return batchLoadChannelData(
+			ctx, cfg.QueryCfg, db, channelIDs, policyIDs,
+		)
+	}
+
+	processItem := func(ctx context.Context,
+		row sqlc.ListChannelsWithPoliciesPaginatedRow,
+		batchData *batchChannelData) error {
+
+		node1, node2, err := buildNodeVertices(
+			row.Node1Pubkey, row.Node2Pubkey,
+		)
+		if err != nil {
+			return err
+		}
+
+		edge, err := buildEdgeInfoWithBatchData(
+			cfg.ChainHash, row.GraphChannel, node1, node2,
+			batchData,
+		)
+		if err != nil {
+			return fmt.Errorf("unable to build channel info: %w",
+				err)
+		}
+
+		dbPol1, dbPol2, err := extractChannelPolicies(row)
+		if err != nil {
+			return err
+		}
+
+		p1, p2, err := buildChanPoliciesWithBatchData(
+			dbPol1, dbPol2, edge.ChannelID, node1, node2, batchData,
+		)
+		if err != nil {
+			return err
+		}
+
+		return processChannel(edge, p1, p2)
+	}
+
+	return sqldb.ExecuteCollectAndBatchWithSharedDataQuery(
+		ctx, cfg.QueryCfg, int64(-1), pageQueryFunc, extractPageCursor,
+		collectFunc, batchDataFunc, processItem,
+	)
+}

make/testing_flags.mk

@@ -121,6 +121,12 @@ ifneq ($(nocache),)
 TEST_FLAGS += -test.count=1
 endif
 
+# If the short flag is added, then any unit tests marked with "testing.Short()"
+# will be skipped.
+ifneq ($(short),)
+TEST_FLAGS += -short
+endif
+
 GOLIST := $(GOCC) list -tags="$(DEV_TAGS)" -deps $(PKG)/... | grep '$(PKG)'| grep -v '/vendor/'
 
 # UNIT_TARGTED is undefined iff a specific package and/or unit test case is

sqldb/paginate.go

@@ -157,3 +157,99 @@ func ExecutePaginatedQuery[C any, T any](ctx context.Context, cfg *QueryConfig,
 
 	return nil
 }
+
+// CollectAndBatchDataQueryFunc represents a function that batch loads
+// additional data for collected identifiers, returning the batch data that
+// applies to all items.
+type CollectAndBatchDataQueryFunc[ID any, BatchData any] func(context.Context,
+	[]ID) (BatchData, error)
+
+// ItemWithBatchDataProcessFunc represents a function that processes individual
+// items along with shared batch data.
+type ItemWithBatchDataProcessFunc[T any, BatchData any] func(context.Context,
+	T, BatchData) error
+
+// CollectFunc represents a function that extracts an identifier from a
+// paginated item.
+type CollectFunc[T any, ID any] func(T) (ID, error)
+
+// ExecuteCollectAndBatchWithSharedDataQuery implements a page-by-page
+// processing pattern where each page is immediately processed with batch-loaded
+// data before moving to the next page.
+//
+// It:
+// 1. Fetches a page of items using cursor-based pagination
+// 2. Collects identifiers from that page and batch loads shared data
+// 3. Processes each item in the page with the shared batch data
+// 4. Moves to the next page and repeats
+//
+// Parameters:
+// - initialCursor: starting cursor for pagination
+// - pageQueryFunc: fetches a page of items
+// - extractPageCursor: extracts cursor from paginated item for next page
+// - collectFunc: extracts identifier from paginated item
+// - batchDataFunc: batch loads shared data from collected IDs for one page
+// - processItem: processes each item with the shared batch data
+func ExecuteCollectAndBatchWithSharedDataQuery[C any, T any, I any, D any](
+	ctx context.Context, cfg *QueryConfig, initialCursor C,
+	pageQueryFunc PagedQueryFunc[C, T],
+	extractPageCursor CursorExtractFunc[T, C],
+	collectFunc CollectFunc[T, I],
+	batchDataFunc CollectAndBatchDataQueryFunc[I, D],
+	processItem ItemWithBatchDataProcessFunc[T, D]) error {
+
+	cursor := initialCursor
+
+	for {
+		// Step 1: Fetch the next page of items.
+		items, err := pageQueryFunc(ctx, cursor, cfg.MaxPageSize)
+		if err != nil {
+			return fmt.Errorf("failed to fetch page with "+
+				"cursor %v: %w", cursor, err)
+		}
+
+		// If no items returned, we're done.
+		if len(items) == 0 {
+			break
+		}
+
+		// Step 2: Collect identifiers from this page and batch load
+		// data.
+		pageIDs := make([]I, len(items))
+		for i, item := range items {
+			pageIDs[i], err = collectFunc(item)
+			if err != nil {
+				return fmt.Errorf("failed to collect "+
+					"identifier from item: %w", err)
+			}
+		}
+
+		// Batch load shared data for this page.
+		batchData, err := batchDataFunc(ctx, pageIDs)
+		if err != nil {
+			return fmt.Errorf("failed to load batch data for "+
+				"page: %w", err)
+		}
+
+		// Step 3: Process each item in this page with the shared batch
+		// data.
+		for _, item := range items {
+			err := processItem(ctx, item, batchData)
+			if err != nil {
+				return fmt.Errorf("failed to process item "+
+					"with batch data: %w", err)
+			}
+
+			// Update cursor for next page.
+			cursor = extractPageCursor(item)
+		}
+
+		// If the number of items is less than the max page size,
+		// we assume there are no more items to fetch.
+		if len(items) < int(cfg.MaxPageSize) {
+			break
+		}
+	}
+
+	return nil
+}

sqldb/paginate_test.go

@@ -558,3 +558,533 @@ func TestExecutePaginatedQuery(t *testing.T) {
 		})
 	}
 }
+
+// TestExecuteCollectAndBatchWithSharedDataQuery tests the
+// ExecuteCollectAndBatchWithSharedDataQuery function which processes items in
+// pages, allowing for efficient querying and processing of large datasets with
+// shared data across batches.
+func TestExecuteCollectAndBatchWithSharedDataQuery(t *testing.T) {
+	t.Parallel()
+
+	type channelRow struct {
+		id        int64
+		name      string
+		policyIDs []int64
+	}
+
+	type channelBatchData struct {
+		lookupTable map[int64]string
+		sharedInfo  string
+	}
+
+	type processedChannel struct {
+		id         int64
+		name       string
+		sharedInfo string
+		lookupData string
+	}
+
+	tests := []struct {
+		name                   string
+		maxPageSize            int32
+		allRows                []channelRow
+		initialCursor          int64
+		pageQueryError         error
+		pageQueryErrorOnCall   int
+		batchDataError         error
+		batchDataErrorOnBatch  int
+		processError           error
+		processErrorOnID       int64
+		earlyTerminationOnPage int
+		expectedError          string
+		expectedProcessedItems []processedChannel
+		expectedPageCalls      int
+		expectedBatchCalls     int
+	}{
+		{
+			name:        "multiple pages multiple batches",
+			maxPageSize: 2,
+			allRows: []channelRow{
+				{
+					id:        1,
+					name:      "Chan1",
+					policyIDs: []int64{10, 11},
+				},
+				{
+					id:        2,
+					name:      "Chan2",
+					policyIDs: []int64{20},
+				},
+				{
+					id:        3,
+					name:      "Chan3",
+					policyIDs: []int64{30, 31},
+				},
+				{
+					id:        4,
+					name:      "Chan4",
+					policyIDs: []int64{40},
+				},
+				{
+					id:        5,
+					name:      "Chan5",
+					policyIDs: []int64{50},
+				},
+			},
+			initialCursor: 0,
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+				{
+					id:         2,
+					name:       "Chan2",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-2",
+				},
+				{
+					id:         3,
+					name:       "Chan3",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-3",
+				},
+				{
+					id:         4,
+					name:       "Chan4",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-4",
+				},
+				{
+					id:         5,
+					name:       "Chan5",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-5",
+				},
+			},
+			// Pages: [1,2], [3,4], [5].
+			expectedPageCalls: 3,
+			// One batch call per page with data: [1,2], [3,4], [5].
+			expectedBatchCalls: 3,
+		},
+		{
+			name:          "empty results",
+			maxPageSize:   10,
+			allRows:       []channelRow{},
+			initialCursor: 0,
+			// One call that returns empty.
+			expectedPageCalls: 1,
+			// No batches since no items.
+			expectedBatchCalls: 0,
+		},
+		{
+			name:        "single page single batch",
+			maxPageSize: 10,
+			allRows: []channelRow{
+				{
+					id:        1,
+					name:      "Chan1",
+					policyIDs: []int64{10},
+				},
+				{
+					id:        2,
+					name:      "Chan2",
+					policyIDs: []int64{20},
+				},
+			},
+			initialCursor: 0,
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+				{
+					id:         2,
+					name:       "Chan2",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-2",
+				},
+			},
+			// One page with all items.
+			expectedPageCalls: 1,
+			// One batch call for the single page.
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "page query error first call",
+			maxPageSize: 5,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+			},
+			initialCursor: 0,
+			pageQueryError: errors.New(
+				"database connection failed",
+			),
+			pageQueryErrorOnCall: 1,
+			expectedError: "failed to fetch page with " +
+				"cursor 0",
+			expectedPageCalls:  1,
+			expectedBatchCalls: 0,
+		},
+		{
+			name:        "page query error second call",
+			maxPageSize: 1,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor: 0,
+			pageQueryError: errors.New("database error on " +
+				"second page"),
+			pageQueryErrorOnCall: 2,
+			expectedError: "failed to fetch page with " +
+				"cursor 1",
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+			},
+			expectedPageCalls:  2,
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "batch data error first batch",
+			maxPageSize: 10,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor: 0,
+			batchDataError: errors.New("batch loading " +
+				"failed"),
+			batchDataErrorOnBatch: 1,
+			expectedError: "failed to load batch data " +
+				"for page",
+			expectedPageCalls:  1,
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "batch data error second page",
+			maxPageSize: 1,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor: 0,
+			batchDataError: errors.New("batch loading " +
+				"failed on second page"),
+			batchDataErrorOnBatch: 2,
+			expectedError: "failed to load batch data " +
+				"for page",
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+			},
+			expectedPageCalls:  2,
+			expectedBatchCalls: 2,
+		},
+		{
+			name:        "process error first item",
+			maxPageSize: 10,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor:    0,
+			processError:     errors.New("processing failed"),
+			processErrorOnID: 1,
+			expectedError: "failed to process item with " +
+				"batch data",
+			expectedPageCalls:  1,
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "process error second item",
+			maxPageSize: 10,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor:    0,
+			processError:     errors.New("processing failed"),
+			processErrorOnID: 2,
+			expectedError: "failed to process item with batch " +
+				"data",
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+			},
+			expectedPageCalls:  1,
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "early termination partial page",
+			maxPageSize: 3,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+			},
+			initialCursor:          0,
+			earlyTerminationOnPage: 1,
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         1,
+					name:       "Chan1",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-1",
+				},
+				{
+					id:         2,
+					name:       "Chan2",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-2",
+				},
+			},
+			expectedPageCalls:  1,
+			expectedBatchCalls: 1,
+		},
+		{
+			name:        "different initial cursor",
+			maxPageSize: 2,
+			allRows: []channelRow{
+				{
+					id:   1,
+					name: "Chan1",
+				},
+				{
+					id:   2,
+					name: "Chan2",
+				},
+				{
+					id:   3,
+					name: "Chan3",
+				},
+			},
+			initialCursor: 1,
+			expectedProcessedItems: []processedChannel{
+				{
+					id:         2,
+					name:       "Chan2",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-2",
+				},
+				{
+					id:         3,
+					name:       "Chan3",
+					sharedInfo: "batch-shared",
+					lookupData: "lookup-3",
+				},
+			},
+			// [2,3], [].
+			expectedPageCalls: 2,
+			// One batch call for the page with [2,3].
+			expectedBatchCalls: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			cfg := &QueryConfig{
+				MaxPageSize: tt.maxPageSize,
+			}
+
+			var (
+				processedItems []processedChannel
+				pageCallCount  int
+				batchCallCount int
+			)
+
+			pageQueryFunc := func(ctx context.Context, cursor int64,
+				limit int32) ([]channelRow, error) {
+
+				pageCallCount++
+
+				// Return error on specific call if configured.
+				//nolint:ll
+				if tt.pageQueryErrorOnCall > 0 &&
+					pageCallCount == tt.pageQueryErrorOnCall {
+
+					return nil, tt.pageQueryError
+				}
+
+				// Simulate cursor-based pagination.
+				var items []channelRow
+				for _, row := range tt.allRows {
+					if row.id > cursor &&
+						len(items) < int(limit) {
+
+						items = append(items, row)
+					}
+				}
+
+				// Handle early termination test case.
+				//nolint:ll
+				if tt.earlyTerminationOnPage > 0 &&
+					pageCallCount == tt.earlyTerminationOnPage {
+
+					// Return fewer items than maxPageSize
+					// to trigger termination
+					if len(items) >= int(tt.maxPageSize) {
+						items = items[:tt.maxPageSize-1]
+					}
+				}
+
+				return items, nil
+			}
+
+			extractPageCursor := func(row channelRow) int64 {
+				return row.id
+			}
+
+			collectFunc := func(row channelRow) (int64, error) {
+				return row.id, nil
+			}
+
+			batchDataFunc := func(ctx context.Context,
+				ids []int64) (*channelBatchData, error) {
+
+				batchCallCount++
+
+				// Return error on specific batch if configured.
+				//nolint:ll
+				if tt.batchDataErrorOnBatch > 0 &&
+					batchCallCount == tt.batchDataErrorOnBatch {
+
+					return nil, tt.batchDataError
+				}
+
+				// Create mock batch data.
+				lookupTable := make(map[int64]string)
+				for _, id := range ids {
+					lookupTable[id] =
+						fmt.Sprintf("lookup-%d", id)
+				}
+
+				return &channelBatchData{
+					lookupTable: lookupTable,
+					sharedInfo:  "batch-shared",
+				}, nil
+			}
+
+			processItem := func(ctx context.Context, row channelRow,
+				batchData *channelBatchData) error {
+
+				// Return error on specific item if configured.
+				if tt.processErrorOnID > 0 &&
+					row.id == tt.processErrorOnID {
+
+					return tt.processError
+				}
+
+				processedChan := processedChannel{
+					id:         row.id,
+					name:       row.name,
+					sharedInfo: batchData.sharedInfo,
+					lookupData: batchData.
+						lookupTable[row.id],
+				}
+
+				processedItems = append(
+					processedItems, processedChan,
+				)
+				return nil
+			}
+
+			err := ExecuteCollectAndBatchWithSharedDataQuery(
+				ctx, cfg, tt.initialCursor,
+				pageQueryFunc, extractPageCursor, collectFunc,
+				batchDataFunc, processItem,
+			)
+
+			// Check error expectations.
+			if tt.expectedError != "" {
+				require.ErrorContains(t, err, tt.expectedError)
+				if tt.pageQueryError != nil {
+					require.ErrorIs(
+						t, err, tt.pageQueryError,
+					)
+				}
+				if tt.batchDataError != nil {
+					require.ErrorIs(
+						t, err, tt.batchDataError,
+					)
+				}
+				if tt.processError != nil {
+					require.ErrorIs(
+						t, err, tt.processError,
+					)
+				}
+			} else {
+				require.NoError(t, err)
+			}
+
+			// Check processed results.
+			require.Equal(
+				t, tt.expectedProcessedItems, processedItems,
+			)
+
+			// Check call counts.
+			require.Equal(
+				t, tt.expectedPageCalls, pageCallCount,
+			)
+			require.Equal(
+				t, tt.expectedBatchCalls, batchCallCount,
+			)
+		})
+	}
+}