-
Notifications
You must be signed in to change notification settings - Fork 0
/
batchget.go
117 lines (101 loc) · 3.16 KB
/
batchget.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
package dygo
import (
"context"
"sync"
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
"golang.org/x/sync/errgroup"
)
const opBatchGet = "BatchGet"
// BatchGetItem retrieves multiple items from the DynamoDB table in parallel using multiple threads.
// It takes a context and the number of threads as input parameters.
// It returns a slice of maps, where each map represents an item retrieved from the table, along with any error encountered.
//
// Example:
//
// item := new(Item)
// for _, gId := range gIds {
// db.PK(gId).SK(Equal(SK)).AddBatchGetItem(item, true)
// }
//
// output, err := item.BatchGetItem(context.Background(), 10)
func (i *Item) BatchGetItem(ctx context.Context, threadCount int) ([]map[string]types.AttributeValue, error) {
if i.err != nil {
return nil, i.err
}
g, ctx := errgroup.WithContext(ctx)
g.SetLimit(threadCount)
var mu sync.Mutex
output := make([]map[string]types.AttributeValue, 0)
for _, batch := range i.batchData.batchGet {
batch := batch
g.Go(func() error {
return i.fetchBatch(ctx, batch, &output, &mu)
})
}
if err := g.Wait(); err != nil {
return nil, dynamoError().method(opBatchGet).message(err.Error())
}
return output, nil
}
// BatchGetAuthorizedItem fetches multiple items from DynamoDB in parallel, with authorization.
// It uses a specified number of goroutines to perform the batch get operation concurrently.
// The fetched items are stored in the Output struct and returned as a result.
// Use Unmarshal() to unmarshal the results into the provided 'out' object.
//
// Example:
//
// func (d *dataItem) Authorize(ctx context.Context) error {
// return nil
// }
//
// item := new(Item)
//
// for _, gId := range gIds {
// db.PK(gId).SK(Equal(SK)).AddBatchGetItem(item, true)
// }
//
// var data []dataItem
// err = item.BatchGetAuthorizedItem(context.Background(), 10).
// Unmarshal(&data, []string{"room"}).
// Run()
func (i *Item) BatchGetAuthorizedItem(ctx context.Context, threadCount int) *output {
result := newOutput(i, ctx)
if i.err != nil {
result.item.err = i.err
}
g, ctx := errgroup.WithContext(ctx)
g.SetLimit(threadCount)
var mu sync.Mutex
output := make([]map[string]types.AttributeValue, 0)
for _, batch := range i.batchData.batchGet {
batch := batch
g.Go(func() error {
return i.fetchBatch(ctx, batch, &output, &mu)
})
}
if err := g.Wait(); err != nil {
result.item.err = dynamoError().method(opBatchGet).message(err.Error())
}
result.Results = append(result.Results, output...)
return result
}
// fetchBatch fetches items in batches from DynamoDB using BatchGetItem API.
func (i *Item) fetchBatch(ctx context.Context, batch map[string]types.KeysAndAttributes, output *[]map[string]types.AttributeValue, mu *sync.Mutex) error {
input := &dynamodb.BatchGetItemInput{
RequestItems: batch,
}
paginator := newBatchGetItemPaginator(i.c.client, input)
for paginator.hasMorePages() {
page, err := paginator.nextPage(ctx)
if err != nil {
return dynamoError().method(opBatchGet).message(err.Error())
}
mu.Lock()
for _, items := range page.Responses {
*output = append(*output, items...)
}
mu.Unlock()
}
return nil
}