Migrate to aws-sdk-go-v2

.github/workflows/main.yml

@@ -29,8 +29,10 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-go-
 
+      - run: go install go.uber.org/mock/mockgen@latest
+
       # `go test` uses only a high-confidence subset of go vet, so execute also `go vet`
-      - run: go vet ./...
+      - run: make vet
 
       - run: go fmt ./... && git diff --exit-code -- ':!go.sum'
 

.gitignore

@@ -1,3 +1,4 @@
 /CREDITS
 /bin
 /dist
+mocks.go

Makefile

@@ -1,5 +1,6 @@
 NAME := ecsmec
-SRCS := $(shell find . -type f -name '*.go' -not -name '*_test.go')
+SRCS := $(shell find . -type f -name '*.go' -not -name '*_test.go' -not -path './internal/testing/*')
+MOCKS := internal/testing/capacitymock/mocks.go internal/testing/servicemock/mocks.go
 
 all: bin/$(NAME)
 
@@ -8,12 +9,21 @@ bin/$(NAME): $(SRCS)
 
 .PHONY: clean
 clean:
-	rm -rf bin/$(NAME)
+	rm -rf bin/$(NAME) $(MOCKS)
 
 .PHONY: install
 install:
 	go install -ldflags "-s -w -X github.com/abicky/ecsmec/cmd.revision=$(shell git rev-parse --short HEAD)"
 
 .PHONY: test
-test:
+test: $(MOCKS)
 	go test -v ./...
+
+.PHONY: vet
+vet: $(MOCKS)
+	go vet ./...
+
+$(MOCKS): $(SRCS)
+	go generate ./...
+# mockgen doesn't update timestamps if the generated code doesn't change
+	touch $(MOCKS)

cmd/recreateservice.go

@@ -4,7 +4,7 @@ import (
 	"encoding/json"
 	"strings"
 
-	"github.com/aws/aws-sdk-go/service/ecs"
+	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	"github.com/spf13/cobra"
 
 	"github.com/abicky/ecsmec/internal/service"
@@ -64,12 +64,12 @@ func recreateService(cmd *cobra.Command, args []string) error {
 		return newRuntimeError("failed to parse \"overrides\": %w", err)
 	}
 
-	sess, err := newSession()
+	cfg, err := newConfig(cmd.Context())
 	if err != nil {
 		return newRuntimeError("failed to initialize a session: %w", err)
 	}
 
-	if err := service.NewService(ecs.New(sess)).Recreate(cluster, serviceName, overrideDef); err != nil {
+	if err := service.NewService(ecs.NewFromConfig(cfg)).Recreate(cmd.Context(), cluster, serviceName, overrideDef); err != nil {
 		return newRuntimeError("failed to recreate the service: %w", err)
 	}
 	return nil

cmd/reduceclustercapacity.go

@@ -1,15 +1,18 @@
 package cmd
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/service/autoscaling"
-	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ecs"
-	"github.com/aws/aws-sdk-go/service/eventbridge"
-	"github.com/aws/aws-sdk-go/service/sqs"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/ecs"
+	"github.com/aws/aws-sdk-go-v2/service/eventbridge"
+	eventbridgetypes "github.com/aws/aws-sdk-go-v2/service/eventbridge/types"
+	"github.com/aws/aws-sdk-go-v2/service/sqs"
+	sqstypes "github.com/aws/aws-sdk-go-v2/service/sqs/types"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 
@@ -39,7 +42,7 @@ that belong to the auto scaling group or spot fleet request.`,
 
 	cmd.Flags().String("cluster", "default", "The name of the target `CLUSTER`")
 
-	cmd.Flags().Int64("amount", 0, "The amount of the capacity to reduce (required)")
+	cmd.Flags().Int32("amount", 0, "The amount of the capacity to reduce (required)")
 	cmd.MarkFlagRequired("amount")
 
 	reduceClusterCapacityCmd = cmd
@@ -49,7 +52,7 @@ func reduceClusterCapacity(cmd *cobra.Command, args []string) error {
 	id, _ := reduceClusterCapacityCmd.Flags().GetString("spot-fleet-request-id")
 	name, _ := reduceClusterCapacityCmd.Flags().GetString("auto-scaling-group-name")
 	cluster, _ := reduceClusterCapacityCmd.Flags().GetString("cluster")
-	amount, _ := reduceClusterCapacityCmd.Flags().GetInt64("amount")
+	amount, _ := reduceClusterCapacityCmd.Flags().GetInt32("amount")
 
 	if len(id) == 0 && len(name) == 0 {
 		return errors.New("\"spot-fleet-request-id\" or \"auto-scaling-group-name\" is required")
@@ -58,98 +61,98 @@ func reduceClusterCapacity(cmd *cobra.Command, args []string) error {
 		return errors.New("\"amount\" must be greater than 0")
 	}
 
-	sess, err := newSession()
+	cfg, err := newConfig(cmd.Context())
 	if err != nil {
 		return newRuntimeError("failed to initialize a session: %w", err)
 	}
 
-	drainer, err := capacity.NewDrainer(cluster, ecsconst.MaxListableContainerInstances, ecs.New(sess))
+	drainer, err := capacity.NewDrainer(cluster, ecsconst.MaxListableContainerInstances, ecs.NewFromConfig(cfg))
 	if err != nil {
 		return newRuntimeError("failed to initialize a Drainer: %w", err)
 	}
 
 	if len(id) == 0 {
-		asg, err := capacity.NewAutoScalingGroup(name, autoscaling.New(sess), ec2.New(sess))
+		asg, err := capacity.NewAutoScalingGroup(name, autoscaling.NewFromConfig(cfg), ec2.NewFromConfig(cfg))
 		if err != nil {
 			return newRuntimeError("failed to initialize a AutoScalingGroup: %w", err)
 		}
 
-		if err := asg.ReduceCapacity(amount, drainer); err != nil {
+		if err := asg.ReduceCapacity(cmd.Context(), amount, drainer); err != nil {
 			return newRuntimeError("failed to reduce the cluster capacity: %w", err)
 		}
 	} else {
-		sfr, err := capacity.NewSpotFleetRequest(id, ec2.New(sess))
+		sfr, err := capacity.NewSpotFleetRequest(id, ec2.NewFromConfig(cfg))
 		if err != nil {
 			return newRuntimeError("failed to initialize a SpotFleetRequest: %w", err)
 		}
 
-		sqsSvc := sqs.New(sess)
-		queueURL, queueArn, err := putSQSQueue(sqsSvc, queueNameForInterruptionWarnings)
+		sqsSvc := sqs.NewFromConfig(cfg)
+		queueURL, queueArn, err := putSQSQueue(cmd.Context(), sqsSvc, queueNameForInterruptionWarnings)
 		if err != nil {
 			return newRuntimeError("failed to create a queue for interruption warnings: %w", err)
 		}
 
-		eventsSvc := eventbridge.New(sess)
+		eventsSvc := eventbridge.NewFromConfig(cfg)
 		targetID := "sqs"
-		if err := putEventRule(eventsSvc, sqsSvc, ruleNameForInterruptionWarnings, targetID, queueURL, queueArn); err != nil {
+		if err := putEventRule(cmd.Context(), eventsSvc, sqsSvc, ruleNameForInterruptionWarnings, targetID, queueURL, queueArn); err != nil {
 			return newRuntimeError("failed to create an event rule for interruption warnings: %w", err)
 		}
 
-		if err := sfr.ReduceCapacity(amount, drainer, capacity.NewSQSQueuePoller(queueURL, sqsSvc)); err != nil {
+		if err := sfr.ReduceCapacity(cmd.Context(), amount, drainer, capacity.NewSQSQueuePoller(queueURL, sqsSvc)); err != nil {
 			return newRuntimeError("failed to reduce the cluster capacity: %w", err)
 		}
 
-		if err := deleteEventRule(eventsSvc, ruleNameForInterruptionWarnings, targetID); err != nil {
+		if err := deleteEventRule(cmd.Context(), eventsSvc, ruleNameForInterruptionWarnings, targetID); err != nil {
 			return newRuntimeError("failed to delete the event rule \"%s\": %w", ruleNameForInterruptionWarnings, err)
 		}
-		if err := deleteSQSQueue(sqsSvc, queueURL); err != nil {
+		if err := deleteSQSQueue(cmd.Context(), sqsSvc, queueURL); err != nil {
 			return newRuntimeError("failed to delete the SQS queue \"%s\": %w", queueNameForInterruptionWarnings, err)
 		}
 	}
 
 	return nil
 }
 
-func putSQSQueue(svc *sqs.SQS, name string) (string, string, error) {
-	queue, err := svc.CreateQueue(&sqs.CreateQueueInput{
+func putSQSQueue(ctx context.Context, svc *sqs.Client, name string) (string, string, error) {
+	queue, err := svc.CreateQueue(ctx, &sqs.CreateQueueInput{
 		QueueName: aws.String(name),
 	})
 	if err != nil {
 		return "", "", xerrors.Errorf("failed to create the SQS queue \"%s\": %w", name, err)
 	}
 
-	attrs, err := svc.GetQueueAttributes(&sqs.GetQueueAttributesInput{
-		AttributeNames: []*string{
-			aws.String("QueueArn"),
+	attrs, err := svc.GetQueueAttributes(ctx, &sqs.GetQueueAttributesInput{
+		AttributeNames: []sqstypes.QueueAttributeName{
+			"QueueArn",
 		},
 		QueueUrl: queue.QueueUrl,
 	})
 	if err != nil {
 		return "", "", xerrors.Errorf("failed to get queue attributes of the queue \"%s\": %w", name, err)
 	}
 
-	return *queue.QueueUrl, *attrs.Attributes["QueueArn"], nil
+	return *queue.QueueUrl, attrs.Attributes["QueueArn"], nil
 }
 
-func deleteSQSQueue(svc *sqs.SQS, queueURL string) error {
-	_, err := svc.DeleteQueue(&sqs.DeleteQueueInput{
+func deleteSQSQueue(ctx context.Context, svc *sqs.Client, queueURL string) error {
+	_, err := svc.DeleteQueue(ctx, &sqs.DeleteQueueInput{
 		QueueUrl: aws.String(queueURL),
 	})
 	return err
 }
 
-func putEventRule(eventsSvc *eventbridge.EventBridge, sqsSvc *sqs.SQS, ruleName, targetID, queueURL, queueArn string) error {
-	rule, err := eventsSvc.PutRule(&eventbridge.PutRuleInput{
+func putEventRule(ctx context.Context, eventsSvc *eventbridge.Client, sqsSvc *sqs.Client, ruleName, targetID, queueURL, queueArn string) error {
+	rule, err := eventsSvc.PutRule(ctx, &eventbridge.PutRuleInput{
 		EventPattern: aws.String("{\"detail-type\":[\"EC2 Spot Instance Interruption Warning\"],\"source\":[\"aws.ec2\"]}"),
 		Name:         aws.String(ruleName),
 	})
 	if err != nil {
 		return xerrors.Errorf("failed to create a rule for interruption warnings: %w", err)
 	}
 
-	_, err = sqsSvc.SetQueueAttributes(&sqs.SetQueueAttributesInput{
-		Attributes: map[string]*string{
-			"Policy": aws.String(fmt.Sprintf(`{
+	_, err = sqsSvc.SetQueueAttributes(ctx, &sqs.SetQueueAttributesInput{
+		Attributes: map[string]string{
+			"Policy": fmt.Sprintf(`{
               "Version": "2012-10-17",
               "Statement": [
                 {
@@ -166,17 +169,17 @@ func putEventRule(eventsSvc *eventbridge.EventBridge, sqsSvc *sqs.SQS, ruleName,
                   }
                 }
               ]
-		    }`, queueArn, *rule.RuleArn)),
+		    }`, queueArn, *rule.RuleArn),
 		},
 		QueueUrl: aws.String(queueURL),
 	})
 	if err != nil {
 		return xerrors.Errorf("failed to update the queue access policy for interruption warnings: %w", err)
 	}
 
-	_, err = eventsSvc.PutTargets(&eventbridge.PutTargetsInput{
+	_, err = eventsSvc.PutTargets(ctx, &eventbridge.PutTargetsInput{
 		Rule: aws.String(ruleName),
-		Targets: []*eventbridge.Target{
+		Targets: []eventbridgetypes.Target{
 			{
 				Id:  aws.String(targetID),
 				Arn: aws.String(queueArn),
@@ -190,17 +193,17 @@ func putEventRule(eventsSvc *eventbridge.EventBridge, sqsSvc *sqs.SQS, ruleName,
 	return nil
 }
 
-func deleteEventRule(svc *eventbridge.EventBridge, ruleName, targetID string) error {
-	_, err := svc.RemoveTargets(&eventbridge.RemoveTargetsInput{
-		Ids:  []*string{aws.String(targetID)},
+func deleteEventRule(ctx context.Context, svc *eventbridge.Client, ruleName, targetID string) error {
+	_, err := svc.RemoveTargets(ctx, &eventbridge.RemoveTargetsInput{
+		Ids:  []string{targetID},
 		Rule: aws.String(ruleName),
 	})
 	if err != nil {
 		return xerrors.Errorf("failed to remove targets of the rule \"%s\": %w", ruleName, err)
 	}
 
-	_, err = svc.DeleteRule(&eventbridge.DeleteRuleInput{
-		Force: aws.Bool(true),
+	_, err = svc.DeleteRule(ctx, &eventbridge.DeleteRuleInput{
+		Force: true,
 		Name:  aws.String(ruleName),
 	})
 	return err

cmd/replaceautoscalinggroupinstances.go

@@ -1,9 +1,9 @@
 package cmd
 
 import (
-	"github.com/aws/aws-sdk-go/service/autoscaling"
-	"github.com/aws/aws-sdk-go/service/ec2"
-	"github.com/aws/aws-sdk-go/service/ecs"
+	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/ecs"
 	"github.com/spf13/cobra"
 
 	"github.com/abicky/ecsmec/internal/capacity"
@@ -28,32 +28,32 @@ launches new ones.`,
 
 	cmd.Flags().String("cluster", "default", "The name of the target `CLUSTER`")
 
-	cmd.Flags().Int64("batch-size", ecsconst.MaxListableContainerInstances, "The number of instances drained at a once")
+	cmd.Flags().Int32("batch-size", ecsconst.MaxListableContainerInstances, "The number of instances drained at a once")
 
 	replaceAutoScalingGroupInstancesCmd = cmd
 }
 
 func replaceAutoScalingGroupInstances(cmd *cobra.Command, args []string) error {
 	name, _ := replaceAutoScalingGroupInstancesCmd.Flags().GetString("auto-scaling-group-name")
 	cluster, _ := replaceAutoScalingGroupInstancesCmd.Flags().GetString("cluster")
-	batchSize, _ := replaceAutoScalingGroupInstancesCmd.Flags().GetInt64("batch-size")
+	batchSize, _ := replaceAutoScalingGroupInstancesCmd.Flags().GetInt32("batch-size")
 
-	sess, err := newSession()
+	cfg, err := newConfig(cmd.Context())
 	if err != nil {
 		return newRuntimeError("failed to initialize a session: %w", err)
 	}
 
-	asg, err := capacity.NewAutoScalingGroup(name, autoscaling.New(sess), ec2.New(sess))
+	asg, err := capacity.NewAutoScalingGroup(name, autoscaling.NewFromConfig(cfg), ec2.NewFromConfig(cfg))
 	if err != nil {
 		return newRuntimeError("failed to initialize a AutoScalingGroup: %w", err)
 	}
 
-	drainer, err := capacity.NewDrainer(cluster, batchSize, ecs.New(sess))
+	drainer, err := capacity.NewDrainer(cluster, batchSize, ecs.NewFromConfig(cfg))
 	if err != nil {
 		return newRuntimeError("failed to initialize a Drainer: %w", err)
 	}
 
-	if err := asg.ReplaceInstances(drainer); err != nil {
+	if err := asg.ReplaceInstances(cmd.Context(), drainer); err != nil {
 		return newRuntimeError("failed to replace instances: %w", err)
 	}
 	return nil

cmd/root.go

@@ -1,10 +1,11 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/spf13/cobra"
 	"golang.org/x/xerrors"
 )
@@ -62,13 +63,8 @@ func init() {
 	rootCmd.PersistentFlags().String("region", "", "The AWS region")
 }
 
-func newSession() (*session.Session, error) {
+func newConfig(ctx context.Context) (aws.Config, error) {
 	region, _ := rootCmd.Flags().GetString("region")
 	profile, _ := rootCmd.Flags().GetString("profile")
-	return session.NewSessionWithOptions(session.Options{
-		Config: aws.Config{
-			Region: aws.String(region),
-		},
-		Profile: profile,
-	})
+	return config.LoadDefaultConfig(ctx, config.WithRegion(region), config.WithSharedConfigProfile(profile))
 }