Full example for DynamoDB batch-get-item in golang

gistfile1.txt

module github.com/tmtk75/m

go 1.16

require (
	github.com/aws/aws-sdk-go-v2/config v1.4.1 // indirect
	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.4.0 // indirect
)

gistfile2.txt

github.com/aws/aws-sdk-go-v2 v1.7.0 h1:UYGnoIPIzed+ycmgw8Snb/0HK+KlMD+SndLTneG8ncE=
github.com/aws/aws-sdk-go-v2 v1.7.0/go.mod h1:tb9wi5s61kTDA5qCkcDbt3KRVV74GGslQkl/DRdX/P4=
github.com/aws/aws-sdk-go-v2/config v1.4.1 h1:PcGp9Kf+1dHJmP3EIDZJmAmWfGABFTU0obuvYQNzWH8=
github.com/aws/aws-sdk-go-v2/config v1.4.1/go.mod h1:HCDWZ/oeY59TPtXslxlbkCqLQBsVu6b09kiG43tdP+I=
github.com/aws/aws-sdk-go-v2/credentials v1.3.0 h1:vXxTINCsHn6LKhR043jwSLd6CsL7KOEU7b1woMr1K1A=
github.com/aws/aws-sdk-go-v2/credentials v1.3.0/go.mod h1:tOcv+qDZ0O+6Jk2beMl5JnZX6N0H7O8fw9UsD3bP7GI=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.2.0 h1:ucExzYCoAiL9GpKOsKkQLsa43wTT23tcdP4cDTSbZqY=
github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.2.0/go.mod h1:XvzoGzuS0kKPzCQtJCC22Xh/mMgVAzfGo/0V+mk/Cu0=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.0.0 h1:A9b5Mvsb4SEZuNzkxH4cenBckc0YgsPncesEmvY8M1I=
github.com/aws/aws-sdk-go-v2/internal/configsources v1.0.0/go.mod h1:gVCp/Wo3eyri2BAFHafQwkpDSldgAyE+4TCGS5zrM44=
github.com/aws/aws-sdk-go-v2/internal/ini v1.1.0 h1:DJq/vXXF+LAFaa/kQX9C6arlf4xX4uaaqGWIyAKOCpM=
github.com/aws/aws-sdk-go-v2/internal/ini v1.1.0/go.mod h1:qGQ/9IfkZonRNSNLE99/yBJ7EPA/h8jlWEqtJCcaj+Q=
github.com/aws/aws-sdk-go-v2/service/dynamodb v1.4.0 h1:EUl9GxhdKy7aqg8cZqCZ5cy/tmYtw/83rZIkmcWVFik=
github.com/aws/aws-sdk-go-v2/service/dynamodb v1.4.0/go.mod h1:M8xNNEkA5Y2wVlXwxToKJEDRNbKzvcWHYao+2XeszIY=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.2.0 h1:wfI4yrOCMAGdHaEreQ65ycSmPLVc2Q82O+r7ZxYTynA=
github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.2.0/go.mod h1:2Kc2Pybp1Hr2ZCCOz78mWnNSZYEKKBQgNcizVGk9sko=
github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.0.0 h1:cHNurcHJYuifPyTHJe2NTOXbyXw5Ga2hMNrazCfJvjM=
github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.0.0/go.mod h1:+66FlGqa06oNFJImpstlS7CFrGIL+lvSobhkY19ukSw=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.2.0 h1:g2npzssI/6XsoQaPYCxliMFeC5iNKKvO0aC+/wWOE0A=
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.2.0/go.mod h1:a7XLWNKuVgOxjssEF019IiHPv35k8KHBaWv/wJAfi2A=
github.com/aws/aws-sdk-go-v2/service/sso v1.3.0 h1:DMi9w+TpUam7eJ8ksL7svfzpqpqem2MkDAJKW8+I2/k=
github.com/aws/aws-sdk-go-v2/service/sso v1.3.0/go.mod h1:qWR+TUuvfji9udM79e4CPe87C5+SjMEb2TFXkZaI0Vc=
github.com/aws/aws-sdk-go-v2/service/sts v1.5.0 h1:Y1K9dHE2CYOWOvaJSIITq4mJfLX43iziThTvqs5FqOg=
github.com/aws/aws-sdk-go-v2/service/sts v1.5.0/go.mod h1:HjDKUmissf6Mlut+WzG2r35r6LeTKmLEDJ6p9NryzLg=
github.com/aws/smithy-go v1.5.0 h1:2grDq7LxZlo8BZUDeqRfQnQWLZpInmh2TLPPkJku3YM=
github.com/aws/smithy-go v1.5.0/go.mod h1:SObp3lf9smib00L/v3U2eAKG8FyQ7iLrJnQiAmR5n+E=
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

main.go

package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
)

var logger = log.Default()

func main() {
	tn := flag.String("table-name", "", "a DynamoDB table name")
	N := flag.Int("batch-size", 100, "size of batch items")
	kn := flag.String("key-name", "", "primary key name")
	st := flag.Duration("sleep", 10*time.Millisecond, "wait to send next request")
	to := flag.Duration("timeout", 30*time.Second, "timeout of context")
	flag.Parse()

	logger.Printf("table-name: %v", *tn)
	logger.Printf("key-name: %v", *kn)

	ctx := context.Background()
	ctx, cancel := context.WithTimeout(context.Background(), *to)
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, os.Interrupt)
	go func() {
		select {
		case <-ch:
			cancel()
		}
	}()

	cfg, err := config.LoadDefaultConfig(ctx)
	if err != nil {
		panic("unable to load SDK config, " + err.Error())
	}

	svc := dynamodb.NewFromConfig(cfg)
	if err != nil {
		log.Fatalf("%v", err)
	}

	type Keys []map[string]types.AttributeValue
	allKeys := make([]Keys, 0)
	keys := make(Keys, 0)

	ln := 0
	ri := bufio.NewScanner(os.Stdin)
	for ri.Scan() {
		l := ri.Text()
		if l == "" {
			break
		}
		ln += 1

		keys = append(keys, map[string]types.AttributeValue{*kn: &types.AttributeValueMemberS{Value: l}})
		if len(keys) == *N {
			allKeys = append(allKeys, keys)
			keys = make(Keys, 0)
		}
	}
	if len(keys) > 0 {
		allKeys = append(allKeys, keys)
	}
	if len(allKeys) == 0 {
		log.Fatalf("must have length greater than or eaual to 1. Check your stdin")
	}
	logger.Printf("read %v lines", ln)

	rch := make(chan func() *dynamodb.BatchGetItemOutput)
	var wg sync.WaitGroup
	for _, k := range allKeys {
		go func(k Keys) {
			rch <- func() *dynamodb.BatchGetItemOutput {
				return func(k Keys) *dynamodb.BatchGetItemOutput {
					defer wg.Done()
					r, err := svc.BatchGetItem(ctx, &dynamodb.BatchGetItemInput{
						// https://docs.aws.amazon.com/ja_jp/amazondynamodb/latest/APIReference/API_BatchGetItem.html#API_BatchGetItem_RequestSyntax
						RequestItems: map[string]types.KeysAndAttributes{
							*tn: types.KeysAndAttributes{Keys: k},
						},
					})
					if err != nil {
						log.Fatalf("failed to query, %v", err)
					}
					return r
				}(k)
			}
		}(k)
		wg.Add(1)
	}
	logger.Printf("ready to send requests")

	all := make([]map[string]types.AttributeValue, 0)
	go func() {
		m := sync.Mutex{}
		for {
			select {
			case f := <-rch:
				go func() {
					defer m.Unlock()
					r := f()
					m.Lock()
					all = append(all, r.Responses[*tn]...)

				}()
				break
			default:
				break
			}
			time.Sleep(*st)
		}
	}()
	wg.Wait()

	b, err := json.Marshal(all)
	if err != nil {
		log.Fatalf("failed to marshal, %v", err)
	}

	fmt.Printf("%v", string(b))
}