Skip to content

Commit c014ee4

Browse files
Santosh1608Santosh1608
authored
feat: Add BatchPutItems method to WriteAPI (#31)
Add BatchPutItems method to WriteAPI So that package users can make use of BatchPutItems exposed by dynago helper pkg Tested locally and it works correctly with these changes
1 parent dd62f46 commit c014ee4

File tree

3 files changed

+240
-14
lines changed

3 files changed

+240
-14
lines changed

‎batch_write_items.go‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,30 +2,48 @@ package dynago
22

33
import (
44
"context"
5+
"fmt"
56

7+
"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
68
"github.com/aws/aws-sdk-go-v2/service/dynamodb"
79
"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
810
)
911

1012
const ChunkSize = 25
1113

14+
type BatchPutItemsInput struct {
15+
PartitionKeyValue Attribute
16+
SortKeyValue Attribute
17+
Item any
18+
}
19+
1220
/**
13-
* Used to update records to dynamodb
14-
* @param input slice of record want to put to DB
15-
* @return error
21+
* BatchPutItems writes multiple items to DynamoDB in batches.
22+
* Items are automatically chunked into groups of 25 (DynamoDB's batch limit).
23+
* Each item is marshaled and partition/sort keys are added before writing.
24+
* @param ctx context for the operation
25+
* @param inputs slice of items to put into DynamoDB
26+
* @return error if operation fails
1627
*/
17-
18-
func (t *Client) BatchWriteItems(ctx context.Context, input []map[string]types.AttributeValue) error {
19-
items := make([]types.WriteRequest, 0, len(input))
28+
func (t *Client) BatchPutItems(ctx context.Context, inputs []BatchPutItemsInput) error {
29+
items := make([]types.WriteRequest, 0, len(inputs))
2030
table := t.TableName
21-
for _, model := range input {
22-
items = append(items,
23-
types.WriteRequest{
24-
PutRequest: &types.PutRequest{
25-
Item: model,
26-
},
31+
32+
for _, input := range inputs {
33+
item, err := attributevalue.MarshalMap(input.Item)
34+
if err != nil {
35+
return fmt.Errorf("failed to marshall item; %s", err)
36+
}
37+
38+
for k, v := range t.NewKeys(input.PartitionKeyValue, input.SortKeyValue) {
39+
item[k] = v
40+
}
41+
42+
items = append(items, types.WriteRequest{
43+
PutRequest: &types.PutRequest{
44+
Item: item,
2745
},
28-
)
46+
})
2947
}
3048
chunkedItems := chunkBy(items, ChunkSize)
3149
for _, chunkedBatch := range chunkedItems {
@@ -40,5 +58,4 @@ func (t *Client) BatchWriteItems(ctx context.Context, input []map[string]types.A
4058
}
4159

4260
return nil
43-
4461
}

‎interface.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ type WriteAPI interface {
2727
PutItem(ctx context.Context, pk, sk Attribute, item interface{}, opt ...PutOption) error
2828
DeleteItem(ctx context.Context, pk, sk string) error
2929
BatchDeleteItems(ctx context.Context, input []AttributeRecord) []AttributeRecord
30+
BatchPutItems(ctx context.Context, items []BatchPutItemsInput) error
3031
}
3132

3233
type TransactionAPI interface {

‎tests/batch_put_items_test.go‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
package tests
2+
3+
import (
4+
"context"
5+
"reflect"
6+
"testing"
7+
8+
"github.com/oolio-group/dynago"
9+
)
10+
11+
type BatchTestRecord struct {
12+
ID string
13+
Name string
14+
Age int
15+
}
16+
17+
func TestBatchPutItems(t *testing.T) {
18+
table := prepareTable(t)
19+
ctx := context.TODO()
20+
21+
testCases := []struct {
22+
title string
23+
inputs []dynago.BatchPutItemsInput
24+
expected []BatchTestRecord
25+
}{
26+
{
27+
title: "single item batch",
28+
inputs: []dynago.BatchPutItemsInput{
29+
{
30+
PartitionKeyValue: dynago.StringValue("user_1"),
31+
SortKeyValue: dynago.StringValue("profile_1"),
32+
Item: BatchTestRecord{
33+
ID: "1",
34+
Name: "John Doe",
35+
Age: 30,
36+
},
37+
},
38+
},
39+
expected: []BatchTestRecord{
40+
{
41+
ID: "1",
42+
Name: "John Doe",
43+
Age: 30,
44+
},
45+
},
46+
},
47+
{
48+
title: "multiple items batch",
49+
inputs: []dynago.BatchPutItemsInput{
50+
{
51+
PartitionKeyValue: dynago.StringValue("user_1"),
52+
SortKeyValue: dynago.StringValue("profile_1"),
53+
Item: BatchTestRecord{
54+
ID: "1",
55+
Name: "John Doe",
56+
Age: 30,
57+
},
58+
},
59+
{
60+
PartitionKeyValue: dynago.StringValue("user_2"),
61+
SortKeyValue: dynago.StringValue("profile_2"),
62+
Item: BatchTestRecord{
63+
ID: "2",
64+
Name: "Jane Smith",
65+
Age: 25,
66+
},
67+
},
68+
{
69+
PartitionKeyValue: dynago.StringValue("user_3"),
70+
SortKeyValue: dynago.StringValue("profile_3"),
71+
Item: BatchTestRecord{
72+
ID: "3",
73+
Name: "Bob Johnson",
74+
Age: 35,
75+
},
76+
},
77+
},
78+
expected: []BatchTestRecord{
79+
{
80+
ID: "1",
81+
Name: "John Doe",
82+
Age: 30,
83+
},
84+
{
85+
ID: "2",
86+
Name: "Jane Smith",
87+
Age: 25,
88+
},
89+
{
90+
ID: "3",
91+
Name: "Bob Johnson",
92+
Age: 35,
93+
},
94+
},
95+
},
96+
{
97+
title: "batch with more than 25 items (chunking test)",
98+
inputs: func() []dynago.BatchPutItemsInput {
99+
var inputs []dynago.BatchPutItemsInput
100+
for i := 1; i <= 30; i++ {
101+
inputs = append(inputs, dynago.BatchPutItemsInput{
102+
PartitionKeyValue: dynago.StringValue("batch_test"),
103+
SortKeyValue: dynago.StringValue("item_" + string(rune(i+'0'))),
104+
Item: BatchTestRecord{
105+
ID: string(rune(i + '0')),
106+
Name: "User " + string(rune(i+'0')),
107+
Age: 20 + i,
108+
},
109+
})
110+
}
111+
return inputs
112+
}(),
113+
expected: func() []BatchTestRecord {
114+
var expected []BatchTestRecord
115+
for i := 1; i <= 30; i++ {
116+
expected = append(expected, BatchTestRecord{
117+
ID: string(rune(i + '0')),
118+
Name: "User " + string(rune(i+'0')),
119+
Age: 20 + i,
120+
})
121+
}
122+
return expected
123+
}(),
124+
},
125+
}
126+
127+
for _, tc := range testCases {
128+
t.Run(tc.title, func(t *testing.T) {
129+
t.Parallel()
130+
131+
err := table.BatchPutItems(ctx, tc.inputs)
132+
if err != nil {
133+
t.Fatalf("unexpected error: %s", err)
134+
}
135+
136+
// Verify all items were written correctly
137+
for i, input := range tc.inputs {
138+
var result BatchTestRecord
139+
err, found := table.GetItem(ctx, input.PartitionKeyValue, input.SortKeyValue, &result)
140+
if err != nil {
141+
t.Fatalf("unexpected error retrieving item %d: %s", i, err)
142+
}
143+
if !found {
144+
t.Fatalf("item %d not found", i)
145+
}
146+
if !reflect.DeepEqual(tc.expected[i], result) {
147+
t.Errorf("item %d: expected %+v, got %+v", i, tc.expected[i], result)
148+
}
149+
}
150+
})
151+
}
152+
}
153+
154+
func TestBatchPutItemsError(t *testing.T) {
155+
table := prepareTable(t)
156+
ctx := context.TODO()
157+
158+
testCases := []struct {
159+
title string
160+
inputs []dynago.BatchPutItemsInput
161+
expectError bool
162+
}{
163+
{
164+
title: "empty partition key should error",
165+
inputs: []dynago.BatchPutItemsInput{
166+
{
167+
PartitionKeyValue: dynago.StringValue(""),
168+
SortKeyValue: dynago.StringValue("profile_1"),
169+
Item: BatchTestRecord{
170+
ID: "1",
171+
Name: "John Doe",
172+
Age: 30,
173+
},
174+
},
175+
},
176+
expectError: true,
177+
},
178+
{
179+
title: "empty sort key should error",
180+
inputs: []dynago.BatchPutItemsInput{
181+
{
182+
PartitionKeyValue: dynago.StringValue("user_1"),
183+
SortKeyValue: dynago.StringValue(""),
184+
Item: BatchTestRecord{
185+
ID: "1",
186+
Name: "John Doe",
187+
Age: 30,
188+
},
189+
},
190+
},
191+
expectError: true,
192+
},
193+
}
194+
195+
for _, tc := range testCases {
196+
t.Run(tc.title, func(t *testing.T) {
197+
t.Parallel()
198+
199+
err := table.BatchPutItems(ctx, tc.inputs)
200+
if tc.expectError && err == nil {
201+
t.Fatalf("expected error but got none")
202+
}
203+
if !tc.expectError && err != nil {
204+
t.Fatalf("unexpected error: %s", err)
205+
}
206+
})
207+
}
208+
}

0 commit comments

Comments
 (0)