Merge pull request #10121 from ellemouton/graphPerf6
[5] sqldb+graph/db: add and use new paginate & batch helper
This commit is contained in:
6 files changed
+818
-149
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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) {
|
||||
|
||||
+177
-149
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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,
|
||||
)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user