From 1a60edbd335e5d9da771c885f5f9c95d1a8053df Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Thu, 31 Jul 2025 10:55:23 +0200 Subject: [PATCH 1/4] make+graph/db: add -short helper and skip tests in short mode In this commit, we add a makefile helper that can be used to add the "-short" flag to the "go test" command when running unit tests via "make unit ...". Tests can then be expanded with a `testing.Short()` check to potentially skip the test in short mode. This is useful for if a dev wants to quickly run most of the tests in a package but would like to opt out of running the longer form tests such as stress tests or rapid generation tests. This commit adds this check to two graph/db tests. --- graph/db/graph_test.go | 5 +++++ graph/db/sql_migration_test.go | 4 ++++ make/testing_flags.mk | 6 ++++++ 3 files changed, 15 insertions(+) diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go index 25ce1edf7..3c60826e1 100644 --- a/graph/db/graph_test.go +++ b/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) diff --git a/graph/db/sql_migration_test.go b/graph/db/sql_migration_test.go index eaaaee6c5..7d6b0e3b0 100644 --- a/graph/db/sql_migration_test.go +++ b/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) { diff --git a/make/testing_flags.mk b/make/testing_flags.mk index 7e4ba5b0e..95a493930 100644 --- a/make/testing_flags.mk +++ b/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 From 905941067e7480a96f3eb49a4cc5bccdb9c90156 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Thu, 31 Jul 2025 12:39:13 +0200 Subject: [PATCH 2/4] sqldb: add ExecuteCollectAndBatchWithSharedDataQuery helper In this commit we add a new helper method in the sqldb package: ExecuteCollectAndBatchWithSharedDataQuery. This can be used to paginate through items in the database while at the same time performing batch data collection for those items. --- sqldb/paginate.go | 96 ++++++++ sqldb/paginate_test.go | 530 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 626 insertions(+) diff --git a/sqldb/paginate.go b/sqldb/paginate.go index 4d8caddb6..2a51218c5 100644 --- a/sqldb/paginate.go +++ b/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 +} diff --git a/sqldb/paginate_test.go b/sqldb/paginate_test.go index 1de3a02b8..dff0a16b1 100644 --- a/sqldb/paginate_test.go +++ b/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, + ) + }) + } +} From 1219cdb7f14f0522d749fa23052204e5d65b50fa Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Thu, 31 Jul 2025 12:43:10 +0200 Subject: [PATCH 3/4] graph/db: use sqldb helper for ForEachNode A pure refactor commit which updates the ForEachNode method to make use of the new sqldb.ExecuteCollectAndBatchWithSharedDataQuery helper. --- graph/db/sql_store.go | 103 +++++++++++++++++++++++++----------------- 1 file changed, 62 insertions(+), 41 deletions(-) diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go index c7366e155..886d1ddcc 100644 --- a/graph/db/sql_store.go +++ b/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 { - err := cb(newSQLGraphNodeTx( - db, s.cfg.ChainHash, dbID, node, - )) - if err != nil { - return fmt.Errorf("callback failed for "+ - "node(id=%d): %w", dbID, err) - } - lastID = dbID + return forEachNodePaginated( + ctx, s.cfg.QueryCfg, db, + ProtocolV1, + func(ctx context.Context, dbNodeID int64, + node *models.LightningNode) error { - return nil - } - - 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 + return cb(newSQLGraphNodeTx( + db, s.cfg.ChainHash, dbNodeID, node, + )) + }, + ) }, reset) } @@ -5082,3 +5051,55 @@ 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, + ) +} From ae13158b686cf529b99bcbda07ab5e11adcfaef0 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Thu, 31 Jul 2025 12:54:13 +0200 Subject: [PATCH 4/4] graph/db: update ForEachChannel to use new sqldb helper Refactor to let ForEachChannel make use of the new sqldb.ExecuteCollectAndBatchWithSharedDataQuery helper. --- graph/db/sql_store.go | 223 ++++++++++++++++++++++-------------------- 1 file changed, 115 insertions(+), 108 deletions(-) diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go index 886d1ddcc..02a3f6372 100644 --- a/graph/db/sql_store.go +++ b/graph/db/sql_store.go @@ -1297,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) - 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 + return forEachChannelWithPolicies(ctx, db, s.cfg, cb) }, reset) } @@ -5103,3 +4996,117 @@ func forEachNodePaginated(ctx context.Context, cfg *sqldb.QueryConfig, 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, + ) +}