dynamodb batch get items sample code

sample.go

const maxRetryCount = 3

func (s *Server) getAlbums(ctx context.Context, albumIDs []uuid.UUID) ([]Album, error) {
	ctx, span := tracer.Start(ctx, "getAlbums")
	defer span.End()

	keysToGet := make([]map[string]types.AttributeValue, len(albumIDs))
	for i, albumID := range albumIDs {
		idv, err := attributevalue.Marshal(albumID.String())
		if err != nil {
			return nil, err
		}
		keysToGet[i] = map[string]types.AttributeValue{
			"AlbumID": idv,
		}
	}

	input := dynamodb.BatchGetItemInput{
		RequestItems: map[string]types.KeysAndAttributes{
			AlbumTableName: {
				Keys: keysToGet,
			},
		},
	}

	ret := make([]Album, 0)

	// retry until UnprocessedKeys is null or reach maxRetryCount 
	for i := 0; i < maxRetryCount; i++ {
		response, err := s.c.BatchGetItem(ctx, &input)
		if err != nil {
			return nil, err
		}

		var resValue []Album

		v, ok := response.Responses[AlbumTableName]
		if !ok {
			if len(response.UnprocessedKeys) == 0 {
				return nil, fmt.Errorf("no AlbumTableName table in response")
			}
			continue // retry
		}
		if err := attributevalue.UnmarshalListOfMaps(v, &resValue); err != nil {
			return nil, err
		}

		if len(response.UnprocessedKeys) == 0 {
			break
		}
		input.RequestItems = response.UnprocessedKeys

		ret = append(ret, resValue...)

		time.Sleep(500 * time.Millisecond)
	}

	return ret, nil
}