Docstore/awsdynamodb: Supporting the AWS V2 SDK

docstore/awsdynamodb/benchmark_test.go

@@ -15,16 +15,15 @@
 package awsdynamodb
 
 import (
+	"context"
 	"fmt"
-	"net/http"
 	"strconv"
 	"testing"
 
-	"github.com/aws/aws-sdk-go/aws"
-	awscreds "github.com/aws/aws-sdk-go/aws/credentials"
-	"github.com/aws/aws-sdk-go/aws/session"
-	"github.com/aws/aws-sdk-go/service/dynamodb"
-	"github.com/aws/aws-sdk-go/service/dynamodb/expression"
+	awsv2cfg "github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression"
+	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
+	dyn2Types "github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
 )
 
 var benchmarkTableName = collectionName3
@@ -35,24 +34,24 @@ func BenchmarkPutVSTransact(b *testing.B) {
 	// The second way calls BatchGetItem followed by TransactWriteItem.
 	//
 	// The results show that separate PutItems are faster for up to two items.
-	sess, err := awsSession(region, http.DefaultClient)
+	cfg, err := awsv2cfg.LoadDefaultConfig(context.Background())
 	if err != nil {
-		b.Fatal(err)
+		b.Fatal("Error initializing aws session for benchmark: ", err)
 	}
-	db := dynamodb.New(sess)
+	db := dynamodb.NewFromConfig(cfg)
 
 	for nItems := 1; nItems <= 5; nItems++ {
 		b.Run(fmt.Sprintf("%d-Items", nItems), func(b *testing.B) {
-			var items []map[string]*dynamodb.AttributeValue
+			var items []map[string]dyn2Types.AttributeValue
 			for i := 0; i < nItems; i++ {
-				items = append(items, map[string]*dynamodb.AttributeValue{
-					"name": new(dynamodb.AttributeValue).SetS(fmt.Sprintf("pt-vs-transact-%d", i)),
-					"x":    new(dynamodb.AttributeValue).SetN(strconv.Itoa(i)),
-					"rev":  new(dynamodb.AttributeValue).SetN("1"),
+				items = append(items, map[string]dyn2Types.AttributeValue{
+					"name": &dyn2Types.AttributeValueMemberS{Value: fmt.Sprintf("pt-vs-transact-%d", i)},
+					"x":    &dyn2Types.AttributeValueMemberN{Value: strconv.Itoa(i)},
+					"rev":  &dyn2Types.AttributeValueMemberN{Value: "1"},
 				})
 			}
 			for _, item := range items {
-				_, err := db.PutItem(&dynamodb.PutItemInput{
+				_, err := db.PutItem(context.Background(), &dynamodb.PutItemInput{
 					TableName: &benchmarkTableName,
 					Item:      item,
 				})
@@ -74,15 +73,15 @@ func BenchmarkPutVSTransact(b *testing.B) {
 	}
 }
 
-func putItems(b *testing.B, db *dynamodb.DynamoDB, items []map[string]*dynamodb.AttributeValue) {
+func putItems(b *testing.B, db *dynamodb.Client, items []map[string]dyn2Types.AttributeValue) {
 	b.Helper()
 
 	for i, item := range items {
-		item["x"].SetN(strconv.Itoa(i + 1))
+		item["x"] = &dyn2Types.AttributeValueMemberN{Value: strconv.Itoa(i + 1)}
 		in := &dynamodb.PutItemInput{
 			TableName:    &benchmarkTableName,
 			Item:         item,
-			ReturnValues: aws.String("ALL_OLD"),
+			ReturnValues: dyn2Types.ReturnValueAllOld,
 		}
 		ce, err := expression.NewBuilder().
 			WithCondition(expression.Name("rev").Equal(expression.Value(1))).
@@ -93,7 +92,7 @@ func putItems(b *testing.B, db *dynamodb.DynamoDB, items []map[string]*dynamodb.
 		in.ExpressionAttributeNames = ce.Names()
 		in.ExpressionAttributeValues = ce.Values()
 		in.ConditionExpression = ce.Condition()
-		out, err := db.PutItem(in)
+		out, err := db.PutItem(context.Background(), in)
 		if err != nil {
 			b.Fatal(err)
 		}
@@ -103,15 +102,15 @@ func putItems(b *testing.B, db *dynamodb.DynamoDB, items []map[string]*dynamodb.
 	}
 }
 
-func batchGetTransactWrite(b *testing.B, db *dynamodb.DynamoDB, items []map[string]*dynamodb.AttributeValue) {
+func batchGetTransactWrite(b *testing.B, db *dynamodb.Client, items []map[string]dyn2Types.AttributeValue) {
 	b.Helper()
 
-	keys := make([]map[string]*dynamodb.AttributeValue, len(items))
-	tws := make([]*dynamodb.TransactWriteItem, len(items))
+	keys := make([]map[string]dyn2Types.AttributeValue, len(items))
+	tws := make([]dyn2Types.TransactWriteItem, len(items))
 	for i, item := range items {
-		keys[i] = map[string]*dynamodb.AttributeValue{"name": items[i]["name"]}
-		item["x"].SetN(strconv.Itoa(i + 2))
-		put := &dynamodb.Put{TableName: &benchmarkTableName, Item: items[i]}
+		keys[i] = map[string]dyn2Types.AttributeValue{"name": items[i]["name"]}
+		item["x"] = &dyn2Types.AttributeValueMemberN{Value: strconv.Itoa(i + 2)}
+		put := &dyn2Types.Put{TableName: &benchmarkTableName, Item: items[i]}
 		ce, err := expression.NewBuilder().
 			WithCondition(expression.Name("rev").Equal(expression.Value(1))).
 			Build()
@@ -121,29 +120,18 @@ func batchGetTransactWrite(b *testing.B, db *dynamodb.DynamoDB, items []map[stri
 		put.ExpressionAttributeNames = ce.Names()
 		put.ExpressionAttributeValues = ce.Values()
 		put.ConditionExpression = ce.Condition()
-		tws[i] = &dynamodb.TransactWriteItem{Put: put}
+		tws[i] = dyn2Types.TransactWriteItem{Put: put}
 	}
-	_, err := db.BatchGetItem(&dynamodb.BatchGetItemInput{
-		RequestItems: map[string]*dynamodb.KeysAndAttributes{
+	_, err := db.BatchGetItem(context.Background(), &dynamodb.BatchGetItemInput{
+		RequestItems: map[string]dyn2Types.KeysAndAttributes{
 			benchmarkTableName: {Keys: keys},
 		},
 	})
 	if err != nil {
 		b.Fatal(err)
 	}
-	_, err = db.TransactWriteItems(&dynamodb.TransactWriteItemsInput{TransactItems: tws})
+	_, err = db.TransactWriteItems(context.Background(), &dynamodb.TransactWriteItemsInput{TransactItems: tws})
 	if err != nil {
 		b.Fatal(err)
 	}
 }
-
-func awsSession(region string, client *http.Client) (*session.Session, error) {
-	// Provide fake creds if running in replay mode.
-	var creds *awscreds.Credentials
-	return session.NewSession(&aws.Config{
-		HTTPClient:  client,
-		Region:      aws.String(region),
-		Credentials: creds,
-		MaxRetries:  aws.Int(0),
-	})
-}